// Do not edit this file for the homework.

// A PBST is a persistent binary search tree.
// PBST and Node objects are immutable: all fields are final.
// PBST uses naive insertion (like BST), so it is not balanced.
//
// When a BST operation would modify the BST, the corresponding PBST
// operation returns a new PBST.  The original PBST remains available
// and unchanged.  The total time and space per modification is O(H),
// where H is the height of the path traversed.

// This is based on share/book/BST.java, but with several changes:
//   * all data fields here (and in class Node) are final
//   * in class Node, the N field is renamed to size
//   * to help with PRedBlackBST, class Node has a 'color' field
//   * mostly removed "private" access (but leaving "public")
//   * rewrote many recursive methods as loops (a bit faster)
//   * frequent use of ternary operator: "<Test> ? <ExpT> : <ExpF>"
//   * no deletion (keeping it simple, for your red-black version)
//   * removed min, max, floor, ceiling (but kept rank, select)
//   * removed linear time keys and range methods
//   * IntIterator, implementing Iterator<Key> using an int and select()
//   * added toString() method (using the iterator)
//   * check() method also checks red-black properties
//   * eliminated dependence on book classes (StdIn/StdOut/Queue)

public class PBST<Key extends Comparable<Key>, Value>
    implements Iterable<Key>
{
    // Fields:
    final Node root;         // the root of our PBST
    final PBST parent;       // our predecessor in history, or null

    // A diagnostic (static==global, not an instance field):
    private static int nodes;   // total nodes created so far
    public static int nodes() { return nodes; }

    // Node is a persistent Node in a PBST (or a PRedBlackBST).
    // In a PBST, the color is always false (no red edges).
    class Node
    {
        final Key key;           // sorted by key
        final Value val;         // associated data
        final Node left, right;  // left and right subtrees
        final boolean color;     // is the link from parent red?
        final int size;          // number of nodes in subtree

        Node(Key k, Value v, Node l, Node r, boolean c) {
            key = k;  val = v;  left = l;  right = r;  color = c;
            size = 1 + size(left) + size(right);
            ++nodes;
        }
        // Methods creating a new Node with one modified field:
        Node setLeft(Node l) {
            return l==left ? this : new Node(key, val, l, right, color);
        }
        Node setRight(Node r) {
            return r==right ? this : new Node(key, val, left, r, color);
        }
        Node setVal(Value v) {
            // For some purposes, .equals would suffice (and save memory).
            return v==val ? this : new Node(key, v, left, right, color);
        }
        Node setColor(boolean c) {
            return c==color ? this : new Node(key, val, left, right, c);
        }
    }

    // Constructors:
    public PBST() { root=null; parent = null; }
    PBST(Node r, PBST p) { root=r; parent=p; assert check(); }

    // Create a new PBST whenever we need a modified root:
    PBST setRoot(Node r) { return r==root ? this : new PBST(r, this); }

    // Node accessor methods, allowing a null Node:
    int size(Node x) { return x==null ? 0 : x.size; }
    Key key(Node x) { return x==null ? null : x.key; }
    Value val(Node x) { return x==null ? null : x.val; }
    boolean isRed(Node x) { return x!=null && x.color; }

    // empty?
    public boolean isEmpty() { return size() == 0; }

    // size of PBST
    public int size() { return size(root); }

    // Does this PBST have the given key?
    public boolean contains(Key key) { return get(root, key) != null; }

    // Return value associated with key, null if no such key.
    public Value get(Key key) { return val(get(root, key)); }

    Node get(Node x, Key key) { // loop, a bit faster than recursion
        while (x != null) {
            int cmp = key.compareTo(x.key);
            if (cmp == 0) return x;
            x = cmp<0 ? x.left : x.right;
        }
        return null;
    }

    // Insert key-value pair in PBST, returning new PBST.
    // May return this, if there is no change.
    public PBST put(Key k, Value v) { return setRoot(put(root, k, v)); }

    // Insert key-value pair in subtree, returning new subtree root.
    // Maybe return the same root, if there is no change.
    Node put(Node x, Key key, Value val) {
        if (x == null) return new Node(key, val, null, null, false);
        int cmp = key.compareTo(x.key);
        if (cmp<0) return x.setLeft(put(x.left, key, val));
        if (cmp>0) return x.setRight(put(x.right, key, val));
        return x.setVal(val);
    }

    // Return key with rank r.  That is, with r smaller keys.
    public Key select(int r) { return key(select(root, r)); }
    Node select(Node x, int r) {
        while (x != null) {
            int t = size(x.left);
            if (r==t) return x;
            if (r<t) { x = x.left; }
            else { x = x.right; r = r-t-1; }
        }
        return x;
    }

    // Return the rank of a key (the number of strictly smaller keys).
    public int rank(Key key) { return rank(key, root); }
    int rank(Key key, Node x) { // in subtree
        int ret = 0;
        while (x != null) {
            int cmp = key.compareTo(x.key);
            if (cmp < 0) { x = x.left; }
            else {
                ret += size(x.left);
                if (cmp==0) break;
                ret += 1;
                x = x.right;
            }
        }
        return ret;
    }

    // Here we implement Iterable<Key>:
    public java.util.Iterator<Key> iterator() { return new IntIterator(); }
    // An IntIterator visits the keys in order.  It uses just an int
    // to keep track of the next key to visit (using select).  In a
    // balanced tree, it needs O(N lg N) time to visit all N keys.
    class IntIterator implements java.util.Iterator<Key> {
        int pos;
        IntIterator() { pos = 0; }
        public Key next() { return select(pos++); }
        public boolean hasNext() { return pos < size(); }
        public void remove() { // we cannot remove
            throw new UnsupportedOperationException();
        }
    }

    // Return string of form "[(key0, val0), (key1, val1), ...]"
    public String toString()
    {
        StringBuilder sb = new StringBuilder();
        sb.append("[");
        // This "enhanced for loop" implicitly uses our iterator:
        for (Key key: this)
            sb.append("(" + key + ", " + get(key) + "), ");
        if (root != null)
            sb.setLength(sb.length()-2); // trim last ", "
        return sb.append("]").toString();
    }

    // Check integrity of this PBST (or PRedBlackBST).
    boolean check()
    {
        boolean good = true;
        if (!isBST(root, null, null))
        { good=false; System.err.println("Not in symmetric order"); }
        if (!isSizeConsistent(root))
        { good=false; System.err.println("Subtree sizes not consistent"); }
        if (!isRankConsistent(root, 0, size()))
        { good=false; System.err.println("Ranks not consistent"); }
        // Check red-black, too, if PRedBlackBST.
        if (this instanceof PRedBlackBST) {
            if (isRed(root))
            { good=false; System.err.println("The root is not black"); }
            if (!is23(root))
            { good=false; System.err.println("Not a 2-3 tree"); }
            int black = 0;
            for (Node x=root; x != null; x = x.right) ++black;
            if (!isBalanced(root, black))
            { good=false; System.err.println("Not balanced"); }
        }
        return good;
    }

    // does this binary tree satisfy strict symmetric order?
    boolean isBST(Node x, Key min, Key max) {
        if (x == null) return true;
        if (min != null && x.key.compareTo(min) <= 0) return false;
        if (max != null && x.key.compareTo(max) >= 0) return false;
        return isBST(x.left, min, x.key) && isBST(x.right, x.key, max);
    }

    // are the size fields correct?
    boolean isSizeConsistent(Node x) {
        if (x == null) return true;
        if (x.size != size(x.left) + size(x.right) + 1) return false;
        return isSizeConsistent(x.left) && isSizeConsistent(x.right);
    }

    // check that rank is consistent with select (avoiding iterator)
    boolean isRankConsistent(Node x, int beg, int end) {
        if (x==null) return beg==end;
        int xrank = beg+size(x.left);
        if (!isRankConsistent(x.left, beg, xrank)) return false;
        if (select(xrank) != x.key) return false;
        if (rank(x.key) != xrank) return false;
        if (!isRankConsistent(x.right, xrank+1, end)) return false;
        return true;
    }

    // Is this a 2-3 tree, each 3-node represented by a left-leaning red edge?
    boolean is23(Node x) {
        if (x == null) return true;
        if (isRed(x.right)) return false;
        if (isRed(x) && isRed(x.left)) // removed "x != root" here
            return false;
        return is23(x.left) && is23(x.right);
    }
    // Is the black depth exactly 'black' in subtree rooted at x?
    private boolean isBalanced(Node x, int black) {
        if (x == null) return black == 0;
        if (!isRed(x)) black--;
        return isBalanced(x.left, black) && isBalanced(x.right, black);
    }

}