// based on: https://algs4.cs.princeton.edu/33balanced/RedBlackBST.java.html

// TODO: implement NavigableSet

// RAM use in a CompressedOOPS JVM (bytes per element):

//   ~12   when elements are inserted in sorted order

//   ~13.7 when elements are inserted in random order

//

// (Obviously more for very small sets.)

sclass HyperCompactTreeSet<A> extends AbstractSet<A> {

  // A symbol table implemented using a left-leaning red-black BST.

  // This is the 2-3 version.

  // Note: We sometimes cast the nullSentinel to A

  // which is technically incorrect but ok because of type erasure.

  // May want to fix just to be clean on the source level too.

  private static final boolean RED   = true;

  private static final boolean BLACK = false;

  // replacement for null elements

  private static final new O nullSentinel;

  private Node<A> root;     // root of the BST

  int size; // size of tree set

  // BST helper node data type

  abstract sclass Node<A> {

    A val;         // associated data

    Node<A> left() { null; } // get left subtree

    abstract Node<A> setLeft(Node<A> left); // set left subtree - return potentially replaced node

    Node<A> right() { null; } // get right subtree

    abstract Node<A> setRight(Node<A> right); // set right subtree - return potentially replaced node

    abstract bool color();

    abstract Node<A> convertToBlack();

    abstract Node<A> convertToRed();

    abstract Node<A> invertColor();

    Node<A> convertToColor(bool color) { ret color == RED ? convertToRed() : convertToBlack(); }

    abstract bool isLeaf();

  }

  // This represents a common case near a red leaf - a black node

  // containing a red leaf in the left slot and null in the right slot.

  // Combined with the direct storage of black leaf children in NonLeaf,

  // this is all we need to get rid of all the leaf overhead. Yay!

  sclass SpecialNode<A> extends Node<A> {

    A leftVal;

    *(A *leftVal, A *val) {}

    bool color() { ret BLACK; }

    Node<A> convertToBlack() { this; }

    Node<A> convertToRed() { ret newNode(RED, val, left(), right()); }

    Node invertColor() { ret convertToRed(); }

    Node<A> left() {

      ret newLeaf(RED, leftVal);

    }

    Node<A> setLeft(Node<A> left) {

      // Can we keep the optimized representation? (Probably this

      // is never going to be true.)

      if (left != null && left.isLeaf() && left.color() == RED)

        ret this with leftVal = left.val;

      else

        ret newNode(BLACK, val, left, right());

    }

    Node<A> right() {

      null;

    }

    Node<A> setRight(Node<A> right) {

      if (right == null) this;

      ret newNode(color(), val, left(), right);

    }

    bool isLeaf() { false; }

  }

  asclass NonLeaf<A> extends Node<A> {

    // either a Node or a (sentinelled) direct user value

    O left, right;

    // color of leaf if left is a user value

    bool defaultLeftLeafColor() { ret BLACK; }

    // color of leaf if right is a user value

    bool defaultRightLeafColor() { ret BLACK; }

    Node<A> left() {

      ret left == null ? null

        : left instanceof Node ? (Node) left

        : newLeaf(defaultLeftLeafColor(), (A) left);

    }

    void setLeft_noMorph(Node<A> left) {

      this.left = left != null && left.isLeaf() && left.color() == defaultLeftLeafColor() ? left.val : left;

    }

    void setRight_noMorph(Node<A> right) {

      this.right = right != null && right.isLeaf() && right.color() == defaultRightLeafColor() ? right.val : right;

    }

    Node<A> setLeft(Node<A> left) {

      ifndef HyperCompactTreeSet_disableSpecialNodes

      if (color() == BLACK && right == null && left != null && left.isLeaf() && left.color() == RED)

        ret new SpecialNode(left.val, val);

      endifndef

      setLeft_noMorph(left);

      if (left == null && right() == null) ret newLeaf(color(), val);

      this;

    }

    Node<A> right() {

      ret right == null ? null

        : right instanceof Node ? (Node) right

        : newLeaf(defaultRightLeafColor(), (A) right);

    }

    Node<A> setRight(Node<A> right) {

      // Setting right to null may produce either a leaf or a

      // special node, so we just go through newNode.

      if (right == null && this.right != null)

        ret newNode(color(), val, left(), null);

      // New right is not null, so we compress (if possible) and store it

      setRight_noMorph(right);

      this;

    }

    bool isLeaf() { false; }

  }

  sclass BlackNode<A> extends NonLeaf<A> {

    *(A *val) {}

    bool color() { ret BLACK; }

    Node<A> convertToBlack() { this; }

    Node<A> convertToRed() { ret newNode(RED, val, left(), right()); }

    Node invertColor() { ret convertToRed(); }

  }

  sclass RedNode<A> extends NonLeaf<A> {

    *(A *val) {}

    bool color() { ret RED; }

    Node<A> convertToBlack() { ret newNode(BLACK, val, left(), right()); }

    Node<A> convertToRed() { this; }

    Node invertColor() { ret convertToBlack(); }

  }

  asclass Leaf<A> extends Node<A> {

    bool isLeaf() { true; }

    Node<A> setLeft(Node<A> left) {

      ret left == null ? this : newNode(color(), val, left, null);

    }

    Node<A> setRight(Node<A> right) {

      ret right == null ? this : newNode(color(), val, null, right);

    }

  }

  sclass BlackLeaf<A> extends Leaf<A> {

    *(A *val) {}

    bool color() { ret BLACK; }

    Node<A> convertToBlack() { this; }

    Node<A> convertToRed() { ret new RedLeaf(val); }

    Node invertColor() { ret convertToRed(); }

  }

  sclass RedLeaf<A> extends Leaf<A> {

    *(A *val) {}

    bool color() { ret RED; }

    Node<A> convertToBlack() { ret new BlackLeaf(val); }

    Node<A> convertToRed() { this; }

    Node invertColor() { ret convertToBlack(); }

  }

  *() {}

  *(Cl<? extends A> cl) { addAll(cl); }

  private static O deSentinel(O o) {

    ret o == nullSentinel ? null : o;

  }

  private static O sentinel(O o) {

    ret o == null ? nullSentinel : o;

  }

  // returns false on null (algorithm needs this)

  static bool <A> isRed(Node<A> x) {

    ret x != null && x.color() == RED;

  }

  static <A> Node<A> newLeaf(bool color, A val) {

    ret color == RED ? new RedLeaf(val) : new BlackLeaf(val);

  }

  static <A> Node<A> newNode(bool color, A val, Node<A> left, Node<A> right) {

    // Make leaf (always a temporary object now)

    if (left == null && right == null)

      ret newLeaf(color, val);

    ifndef HyperCompactTreeSet_disableSpecialNodes

      // Make special node

      if (color == BLACK

        && right == null

        && left != null && left.isLeaf() && left.color() == RED)

        ret new SpecialNode(left.val, val);

    endifndef

    // Make normal non-leaf

    NonLeaf node = color == RED ? new RedNode(val) : new BlackNode(val);

    node.setLeft_noMorph(left);

    node.setRight_noMorph(right);

    ret node;

  }

  public int size() {

    ret size;

  }

  public bool isEmpty() {

    ret root == null;

  }

  public bool add(A val) {

    val = (A) sentinel(val);

    int oldSize = size;

    root = put(root, val);

    root = root.convertToBlack();

    ifdef CompactTreeSet_debug assertTrue(check()); endifdef

    ret size > oldSize;

  }

  // insert the value in the subtree rooted at h

  private Node<A> put(Node<A> h, A val) { 

    if (h == null) { ++size; ret new RedLeaf(val); }

    int cmp = compare_deSentinel(val, h.val);

    if      (cmp < 0) h = h.setLeft(put(h.left(),  val)); 

    else if (cmp > 0) h = h.setRight(put(h.right(), val)) ;

    else              { /*h.val = val;*/ } // no overwriting

    // fix-up any right-leaning links

    if (isRed(h.right()) && !isRed(h.left()))      h = rotateLeft(h);

    if (isRed(h.left())  &&  isRed(h.left().left())) h = rotateRight(h);

    if (isRed(h.left())  &&  isRed(h.right()))     h = flipColors(h);

    ret h;

  }

  final int compare_deSentinel(A a, A b) {

    ret compare((A) deSentinel(a), (A) deSentinel(b));

  }

  // override me if you wish

  int compare(A a, A b) {

    ret cmp(a, b);

  }

  public bool remove(O key) { 

    if (!contains(key)) false;

    key = sentinel(key);

    // if both children of root are black, set root to red

    if (!isRed(root.left()) && !isRed(root.right()))

        root = root.convertToRed();

    root = delete(root, (A) key);

    if (!isEmpty()) root = root.convertToBlack();

    // assert check();

    true;

  }

  // delete the key-value pair with the given key rooted at h

  private Node delete(Node<A> h, A key) { 

    // assert get(h, key) != null;

    if (compare_deSentinel(key, h.val) < 0)  {

      if (!isRed(h.left()) && !isRed(h.left().left()))

          h = moveRedLeft(h);

      h = h.setLeft(delete(h.left(), key));

    }

    else {

      if (isRed(h.left()))

          h = rotateRight(h);

      if (compare_deSentinel(key, h.val) == 0 && (h.right() == null)) {

          --size; null;

      } if (!isRed(h.right()) && !isRed(h.right().left()))

          h = moveRedRight(h);

      if (compare_deSentinel(key, h.val) == 0) {

          --size;

          Node<A> x = min(h.right());

          h.val = x.val;

          // h.val = get(h.right(), min(h.right()).val);

          // h.val = min(h.right()).val;

          h = h.setRight(deleteMin(h.right()));

      }

      else h = h.setRight(delete(h.right(), key));

    }

    return balance(h);

  }

  // make a left-leaning link lean to the right

  private Node<A> rotateRight(Node<A> h) {

      // assert (h != null) && isRed(h.left());

      Node<A> x = h.left();

      h = h.setLeft(x.right());

      x = x.setRight(h);

      x = x.convertToColor(x.right().color());

      x = x.setRight(x.right().convertToRed());

      ret x;

  }

  // make a right-leaning link lean to the left

  private Node<A> rotateLeft(Node<A> h) {

      // assert (h != null) && isRed(h.right());

      Node<A> x = h.right();

      h = h.setRight(x.left());

      x = x.setLeft(h);

      x = x.convertToColor(x.left().color());

      x = x.setLeft(x.left().convertToRed());

      ret x;

  }

  // flip the colors of a node and its two children

  private Node<A> flipColors(Node<A> h) {

      // h must have opposite color of its two children

      // assert (h != null) && (h.left() != null) && (h.right() != null);

      // assert (!isRed(h) &&  isRed(h.left()) &&  isRed(h.right()))

      //    || (isRed(h)  && !isRed(h.left()) && !isRed(h.right()));

      h = h.setLeft(h.left().invertColor());

      h = h.setRight(h.right().invertColor());

      ret h.invertColor();

  }

  // Assuming that h is red and both h.left() and h.left().left()

  // are black, make h.left() or one of its children red.

  private Node<A> moveRedLeft(Node<A> h) {

      // assert (h != null);

      // assert isRed(h) && !isRed(h.left()) && !isRed(h.left().left());

      h = flipColors(h);

      if (isRed(h.right().left())) { 

          h = h.setRight(rotateRight(h.right()));

          h = rotateLeft(h);

          h = flipColors(h);

      }

      ret h;

  }

  // Assuming that h is red and both h.right() and h.right().left()

  // are black, make h.right() or one of its children red.

  private Node<A> moveRedRight(Node<A> h) {

      // assert (h != null);

      // assert isRed(h) && !isRed(h.right()) && !isRed(h.right().left());

      h = flipColors(h);

      if (isRed(h.left().left())) { 

          h = rotateRight(h);

          h = flipColors(h);

      }

      ret h;

  }

  // restore red-black tree invariant

  private Node<A> balance(Node<A> h) {

      // assert (h != null);

      if (isRed(h.right()))                      h = rotateLeft(h);

      if (isRed(h.left()) && isRed(h.left().left())) h = rotateRight(h);

      if (isRed(h.left()) && isRed(h.right()))     h = flipColors(h);

      ret h;

  }

  /**

   * Returns the height of the BST (for debugging).

   * @return the height of the BST (a 1-node tree has height 0)

   */

  public int height() {

      ret height(root);

  }

  private int height(Node<A> x) {

      if (x == null) return -1;

      return 1 + Math.max(height(x.left()), height(x.right()));

  }

  public bool contains(O val) {

    ret find(root, (A) sentinel(val)) != null;

  }

  public A find(A probeVal) {

    probeVal = (A) sentinel(probeVal);

    Node<A> n = find(root, probeVal);

    ret n == null ? null : n.val;

  }

  // value associated with the given key in subtree rooted at x; null if no such key

  private A get(Node<A> x, A key) {

    x = find(x, key);

    ret x == null ? null : x.val;

  }

  Node<A> find(Node<A> x, A key) {

    while (x != null) {

      int cmp = compare_deSentinel(key, x.val);

      if      (cmp < 0) x = x.left();

      else if (cmp > 0) x = x.right();

      else              ret x;

    }

    null;

  }

  private boolean check() {

      if (!is23())             println("Not a 2-3 tree");

      if (!isBalanced())       println("Not balanced");

      return is23() && isBalanced();

  }

  // Does the tree have no red right links, and at most one (left)

  // red links in a row on any path?

  private boolean is23() { return is23(root); }

  private boolean is23(Node<A> x) {

      if (x == null) true;

      if (isRed(x.right())) false;

      if (x != root && isRed(x) && isRed(x.left())) false;

      ret is23(x.left()) && is23(x.right());

  } 

  // do all paths from root to leaf have same number of black edges?

  private bool isBalanced() { 

      int black = 0;     // number of black links on path from root to min

      Node<A> x = root;

      while (x != null) {

          if (!isRed(x)) black++;

          x = x.left();

      }

      ret isBalanced(root, black);

  }

  // does every path from the root to a leaf have the given number of black links?

  private boolean isBalanced(Node<A> x, int black) {

      if (x == null) return black == 0;

      if (!isRed(x)) black--;

      return isBalanced(x.left(), black) && isBalanced(x.right(), black);

  }

  public void clear() { root = null; size = 0; }

  // the smallest key in subtree rooted at x; null if no such key

  private Node<A> min(Node<A> x) { 

    // assert x != null;

    while (x.left() != null) x = x.left();

    ret x;

  }

  private Node<A> deleteMin(Node<A> h) { 

    if (h.left() == null)

      return null;

    if (!isRed(h.left()) && !isRed(h.left().left()))

      h = moveRedLeft(h);

    h = h.setLeft(deleteMin(h.left()));

    ret balance(h);

  }

  public Iterator<A> iterator() {

    ret new MyIterator;

  }

  class MyIterator extends ItIt<A> {

    new L<Node<A>> path;

    *() {

      fetch(root);

    }

    void fetch(Node<A> node) {

      while (node != null) {

        path.add(node);

        node = node.left();

      }

    }

    public bool hasNext() { ret !path.isEmpty(); }

    public A next() {

      if (path.isEmpty()) fail("no more elements");

      Node<A> node = popLast(path);

      // last node is always a leaf, so left is null

      // so proceed to fetch right branch

      fetch(node.right());

      ret (A) deSentinel(node.val);

    }

  }

  // Returns the smallest key in the symbol table greater than or equal to {@code key}.

  public A ceiling(A key) {

    key = (A) sentinel(key);

    Node<A> x = ceiling(root, key);

    ret x == null ? null : x.val;

  }

  // the smallest key in the subtree rooted at x greater than or equal to the given key

  Node<A> ceiling(Node<A> x, A key) {  

    if (x == null) null;

    int cmp = compare_deSentinel(key, x.val);

    if (cmp == 0) ret x;

    if (cmp > 0)  ret ceiling(x.right(), key);

    Node<A> t = ceiling(x.left(), key);

    if (t != null) ret t; 

    else           ret x;

  }

  public A floor(A key) {

    key = (A) sentinel(key);

    Node<A> x = floor(root, key);

    ret x == null ? null : x.val;

  }    

  // the largest key in the subtree rooted at x less than or equal to the given key

  Node<A> floor(Node<A> x, A key) {

    if (x == null) null;

    int cmp = compare_deSentinel(key, x.val);

    if (cmp == 0) ret x;

    if (cmp < 0)  ret floor(x.left(), key);

    Node<A> t = floor(x.right(), key);

    if (t != null) ret t; 

    else           ret x;

  }

  void testInternalStructure() {

    ifndef HyperCompactTreeSet_disableSpecialNodes

      // one leaf object (root) is allowed - could even optimize that

      assertTrue(countLeafObjects() <= 1);

    endifndef

  }

  // count leaf objects we didn't optimize away

  int countLeafObjects(Node node default root) {

    if (node instanceof Leaf) ret 1;

    if (node cast NonLeaf)

      ret countLeafObjects(optCast Node(node.left))

        + countLeafObjects(optCast Node(node.right));

    ret 0;

  }

  Cl<NonLeaf> unoptimizedNodes() {

    new L<NonLeaf> out;

    findUnoptimizedNodes(out);

    ret out;

  }

  void findUnoptimizedNodes(Node node default root, L<NonLeaf> out) {

    if (node == null) ret;

    if (node cast NonLeaf) {

      if (isUnoptimizedNode(node)) out.add(node);

      findUnoptimizedNodes(optCast Node(node.left), out);

      findUnoptimizedNodes(optCast Node(node.right), out);

    }

  }

  bool isUnoptimizedNode(Node node) {

    if (node cast NonLeaf)

      ret node.left instanceof Leaf

        || node.right instanceof Leaf;

    false;

  }

  // Compact me (minimize space use). Returns a compacted copy.

  // This only has an effect when elements were inserted in non-sorted

  // or and/or elements were removed.

  selfType compact() {

    ret new HyperCompactTreeSet(this);

  }

}