// 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 ../book/BST.java but with several changes:
//   * all data fields here are final
//   * field Node.N renamed to Node.size
//   * removed "private" (but left "public" on the intended interface)
//   * rewrote many instances 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 Iterable and range methods
//   * added linear-time toString() method (we also need iterator)
//   * eliminated dependence on book StdIn/StdOut classes

public class PBST<Key extends Comparable<Key>, Value>
{
    final Node root;        // the root of our PBST
    public final int mods;  // number of modifications so far

    // Constructors:
    public PBST() { root=null; mods=0; }
    PBST(Node r, int m) { root=r; mods=m; assert check(); }

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

    class Node
    {
        final Key key;           // sorted by key
        final Value val;         // associated data
        final Node left, right;  // left and right subtrees
        final int size;		 // number of nodes in subtree
	Node(Key k, Value v, Node l, Node r) {
            key = k;
	    val = v;
	    left = l;
	    right = r;
	    size = 1 + size(left) + size(right);
        }
	// Create a new Node when we need a modified field:
	Node setLeft(Node l) {
	    return l==left ? this : new Node(key, val, l, right);
	}
	Node setRight(Node r) {
	    return r==right ? this : new Node(key, val, left, r);
	}
	Node setVal(Value v) {
	    // For some purposes, .equals would be enough.
	    return v==val ? this : new Node(key, v, left, right);
	}
    }
    // 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; }

    // 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);
        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;
    }

    // Produce string of form "[(key1, val1), (key2, val2), ...]"
    public String toString()
    {
        StringBuilder sb = new StringBuilder();
        sb.append("[");
        if (root != null) {
            toString(root, sb);
            sb.setLength(sb.length()-2); // trim the last ", "
        }
        return sb.append("]").toString();
    }
    void toString(Node x, StringBuilder sb)
    {                           // inorder subtree traversal
        if (x==null) return;
        toString(x.left, sb);
        sb.append("("+x.key+", "+x.val+"), ");
        toString(x.right, sb);
    }

    // Check integrity of PBST.
    boolean check() {
	boolean good = true;
        if (!isBST())
	    { good=false; System.err.println("Not in symmetric order"); }
        if (!isSizeConsistent())
	    { good=false; System.err.println("Subtree counts not consistent"); }
        if (!isRankConsistent())
	    { good=false; System.err.println("Ranks not consistent"); }
        return good;
    }

    // does this binary tree satisfy symmetric order?
    boolean isBST() { return isBST(root, null, null); }
    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() { return isSizeConsistent(root); }
    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 ranks are consistent
    boolean isRankConsistent() {
        for (int i = 0; i < size(); i++)
            if (i != rank(select(i))) return false;
        return true;
    }

    // Test client, creating a history of PBST objects.
    public static void main(String[] args) {
        String input = "SEARCHEXAMPLE";
        int N = input.length();
        PBST<String, Integer>[] h = new PBST[N+1];
        h[0] = new PBST();
        System.out.println("h[0] = new PBST();");
        for (int val=0; val<N; ++val) {
            String key = input.substring(val,val+1);
            System.out.printf("h[%d] = h[%d].put(\"%s\", %d);%n",
                               val+1, val, key, val);
            h[val+1] = h[val].put(key, val);
        }
        System.out.printf("%nAll these trees are still usable:%n%n");
        for (int val=0; val <= N; ++val) {
            System.out.printf("h[%d] is %s%n", val, h[val]);
            System.out.printf("h[%d].rank(\"S\") == %d, ", val,
                              h[val].rank("S"));
            System.out.printf("h[%d].get(\"E\") == %s%n", val,
                              h[val].get("E"));
        }
    }
}