// see: https://en.wikipedia.org/wiki/Order_statistic_tree
// based on: https://algs4.cs.princeton.edu/33balanced/RedBlackBST.java.html
final sclass LogNArray<Value> extends RandomAccessAbstractList<Value> {

 // A symbol table implemented using a left-leaning red-black BST.
 // This is the 2-3 version.
 
  private static final boolean RED   = true;
  private static final boolean BLACK = false;

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

  // BST helper node data type
  private sclass Node<Value> {
      private Value val;         // associated data
      private Node<Value> left, right;  // links to left and right subtrees
      private boolean color;     // color of parent link
      private int size;          // subtree count

      public Node(Value val, boolean color, int size) {
          this.val = val;
          this.color = color;
          this.size = size;
      }
  }

  // is node x red; false if x is null ?
  private boolean isRed(Node<Value> x) {
      if (x == null) return false;
      return x.color == RED;
  }

  // number of node in subtree rooted at x; 0 if x is null
  private int size(Node<Value> x) {
      if (x == null) return 0;
      return x.size;
  } 


  public int size() {
      return size(root);
  }

  public boolean isEmpty() {
      return root == null;
  }


  public void add(int idx, Value val) {
    if (idx < 0 || idx > size()) throw new IndexOutOfBoundsException(idx + " / " + size());
    
    root = add(root, idx, val);
    root.color = BLACK;
    ifdef LogNArray_debug assertTrue(check()); endifdef
  }

  // insert the value pair in index k of subtree rooted at h
  private Node<Value> add(Node<Value> h, int k, Value val) { 
    if (h == null) ret new Node(val, RED, 1);

    int t = size(h.left);
    if (k < t)
      // replace / fit in left child
      h.left = add(h.left, k, val);
    else if (k == t) {
      // move value to right child, replace value
      h.right = add(h.right, 0, h.val);
      h.val = val;
    } else
      // replace / fit in right child
      h.right = add(h.right, k-t-1, val); 

    // 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))     flipColors(h);
    h.size = size(h.left) + size(h.right) + 1;
    
    ifdef LogNArray_debug
      print("LogNArray.add: Node " + h.val + " new size: " + h.size);
    endifdef

    ret h;
  }

  public Value remove(int idx) { 
    Value oldValue = get(idx);
    
    // if both children of root are black, set root to red
    if (!isRed(root.left) && !isRed(root.right))
      root.color = RED;

    root = remove(root, idx);
    if (root != null) root.color = BLACK;
    
    ifdef LogNArray_debug assertTrue(check()); endifdef
    
    ret oldValue;
  }

  private Node<Value> remove(Node<Value> h, int k) { 
    int t = size(h.left);
    if (k < t) { // remove from left child
      if (!isRed(h.left) && !isRed(h.left.left))
        h = moveRedLeft(h);
      h.left = remove(h.left, k);
    } else { // remove from center or right child
      if (isRed(h.left)) {
        h = rotateRight(h);
        t = size(h.left);
      }
      if (h.right == null) null; // drop whole node
      if (!isRed(h.right) && !isRed(h.right.left)) {
        h = moveRedRight(h);
        t = size(h.left);
      }
      if (k == t) { // replace center value with min of right child
        h.val = nodeAtIndex(h.right, 0).val;
        h.right = remove(h.right, 0);
      }
      else h.right = remove(h.right, k-t-1);
    }
    ret balance(h);
  }

  // make a left-leaning link lean to the right
  private Node<Value> rotateRight(Node<Value> h) {
      // assert (h != null) && isRed(h.left);
      Node<Value> x = h.left;
      h.left = x.right;
      x.right = h;
      x.color = x.right.color;
      x.right.color = RED;
      x.size = h.size;
      h.size = size(h.left) + size(h.right) + 1;
      return x;
  }

  // make a right-leaning link lean to the left
  private Node<Value> rotateLeft(Node<Value> h) {
      // assert (h != null) && isRed(h.right);
      Node<Value> x = h.right;
      h.right = x.left;
      x.left = h;
      x.color = x.left.color;
      x.left.color = RED;
      x.size = h.size;
      h.size = size(h.left) + size(h.right) + 1;
      return x;
  }

  // flip the colors of a node and its two children
  private void flipColors(Node<Value> 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.color = !h.color;
      h.left.color = !h.left.color;
      h.right.color = !h.right.color;
  }

  // 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<Value> moveRedLeft(Node<Value> h) {
      // assert (h != null);
      // assert isRed(h) && !isRed(h.left) && !isRed(h.left.left);

      flipColors(h);
      if (isRed(h.right.left)) { 
          h.right = rotateRight(h.right);
          h = rotateLeft(h);
          flipColors(h);
      }
      return 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<Value> moveRedRight(Node<Value> h) {
      // assert (h != null);
      // assert isRed(h) && !isRed(h.right) && !isRed(h.right.left);
      flipColors(h);
      if (isRed(h.left.left)) { 
          h = rotateRight(h);
          flipColors(h);
      }
      return h;
  }

  // restore red-black tree invariant
  private Node<Value> balance(Node<Value> 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))     flipColors(h);

      h.size = size(h.left) + size(h.right) + 1;
      return 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() {
      return height(root);
  }
  
  private int height(Node<Value> x) {
      if (x == null) return -1;
      return 1 + Math.max(height(x.left), height(x.right));
  }
  
  public Value get(int k) {
    ret nodeAtIndex(k).val;
  }

  public Value set(int k, Value val) {
    Node<Value> n = nodeAtIndex(k);
    Value oldValue = n.val;
    n.val = val;
    ret oldValue;
  }

  public Node<Value> nodeAtIndex(int k) {
      if (k < 0 || k >= size()) {
          throw new IndexOutOfBoundsException(k + " / " + size());
      }
      ret nodeAtIndex(root, k);
  }

  // the key of rank k in the subtree rooted at x
  private Node<Value> nodeAtIndex(Node<Value> x, int k) {
      // assert x != null;
      // assert k >= 0 && k < size(x);
      int t = size(x.left); 
      if      (t > k) return nodeAtIndex(x.left,  k); 
      else if (t < k) return nodeAtIndex(x.right, k-t-1); 
      else            return x; 
  } 
  
  private boolean check() {
      if (!isSizeConsistent()) println("Subtree counts not consistent");
      if (!is23())             println("Not a 2-3 tree");
      if (!isBalanced())       println("Not balanced");
      return isSizeConsistent() && is23() && isBalanced();
  }

  // are the size fields correct?
  private boolean isSizeConsistent() { return isSizeConsistent(root); }
  private boolean isSizeConsistent(Node<Value> 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);
  } 

  // 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<Value> x) {
      if (x == null) return true;
      if (isRed(x.right)) return false;
      if (x != root && isRed(x) && isRed(x.left))
          return false;
      return is23(x.left) && is23(x.right);
  } 

  // do all paths from root to leaf have same number of black edges?
  private boolean isBalanced() { 
      int black = 0;     // number of black links on path from root to min
      Node<Value> x = root;
      while (x != null) {
          if (!isRed(x)) black++;
          x = x.left;
      }
      return isBalanced(root, black);
  }

  // does every path from the root to a leaf have the given number of black links?
  private boolean isBalanced(Node<Value> 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; }
}

Began life as a copy of #1024412