public class LLRBMap<Key extends Comparable<Key>, Value> {
private static final boolean RED = true;
private static final boolean BLACK = false;
private class Node {
Key key;
Value value;
Node left, right;
int sz, ht;
boolean color;
Node(Key k, Value v, boolean color) {
this.key = k;
this.value = v;
this. color = color;
this.sz = 1;
this.ht = 1;
}
}
private Node root;
private boolean isRed(Node root) {
if (root == null) return false;
return root.color == RED;
}
private int size(Node root) {
if (root == null) return 0;
return root.sz;
}
private void updateSizeAndHeight(Node root) {
root.sz = size(root.left) + size(root.right) + 1;
root.ht = Math.max(height(root.left), height(root.right)) + 1;
}
private Node rotateRight(Node root) {
Node x = root.left;
root.left = x.right;
updateSizeAndHeight(root);
x.right = root;
updateSizeAndHeight(x);
x.color = root.color;
root.color = RED;
return x;
}
private Node rotateLeft(Node root) {
Node x = root.right;
root.right = x.left;
updateSizeAndHeight(root);
x.left = root;
updateSizeAndHeight(x);
x.color = root.color;
root.color = RED;
return x;
}
private void flipColors(Node root) {
root.color = !root.color;
root.left.color = !root.left.color;
root.right.color = !root.right.color;
}
private Node fixUp(Node root) {
if (isRed(root.right) && !isRed(root.left)) root = rotateLeft(root);
if (isRed(root.left) && isRed(root.left.left)) root = rotateRight(root);
if (isRed(root.left) && isRed(root.right)) flipColors(root);
updateSizeAndHeight(root);
return root;
}
private Node put(Node root, Key key, Value value) {
if (root == null) return new Node(key, value, RED);
int cmp = key.compareTo(root.key);
if (cmp < 0) root.left = put(root.left, key, value);
else if (cmp > 0) root.right = put(root.right, key, value);
else root.value = value;
return fixUp(root);
}
private Node moveRedRight(Node root) {
flipColors(root);
if (isRed(root.left.left)) {
root = rotateRight(root);
flipColors(root);
}
return root;
}
private Node moveRedLeft(Node root) {
flipColors(root);
if (isRed(root.right.left)) {
root.right = rotateRight(root.right);
root = rotateLeft(root);
flipColors(root);
}
return root;
}
private Node min(Node root) {
if (root.left == null) return root;
return min(root.left);
}
private Node delete(Node root, Key key) {
if (root == null) return null;
if (key.compareTo(root.key) < 0) {
if (!isRed(root.left) && root.left != null && !isRed(root.left.left)) {
root = moveRedLeft(root);
}
root.left = delete(root.left, key);
}
else {
if (isRed(root.left)) {
root = rotateRight(root);
}
if (key.compareTo(root.key) == 0 && root.right == null) return null;
if (!isRed(root.right) && !isRed(root.right.left)) {
root = moveRedRight(root);
}
if (key.compareTo(root.key) == 0) {
Key tmpKey = root.key;
Node successor = min(root.right);
root.key = successor.key;
root.value = successor.value;
successor.key = tmpKey;
root.right = delete(root.right, tmpKey);
}
else {
root.right = delete(root.right, key);
}
}
return fixUp(root);
}
private Value get(Node root, Key key) {
if (root == null) return null;
int cmp = key.compareTo(root.key);
if (cmp < 0) return get(root.left, key);
else if (cmp > 0) return get(root.right, key);
return root.value;
}
private int rank(Node root, Key key) {
if (root == null) return 0;
int cmp = key.compareTo(root.key);
if (cmp < 0) return rank(root.left, key);
if (cmp == 0) return size(root.left);
return size(root.left) + 1 + rank(root.right, key);
}
private Key select(Node root, int rank) {
if (root == null) return null;
if (size(root.left) == rank) return root.key;
if (size(root.left) > rank) return select(root.left, rank);
return select(root.right, rank - size(root.left) - 1);
}
private int height(Node root) {
if (root == null) return 0;
return root.ht;
}
public void put(Key key, Value value) {
root = put(root, key, value);
}
public Value get(Key key) {
return get(root, key);
}
public void remove(Key key) {
root = delete(root, key);
root.color = BLACK;
}
public int rank(Key key) {
return rank(root, key);
}
public Key select(int rank) {
return select(root, rank);
}
public int height() {
return height(root);
}
private Node copy(Node root) {
if (root == null) return null;
Node x = new Node(root.key, root.value, root.color);
x.ht = root.ht;
x.sz = root.sz;
x.left = copy(root.left);
x.right = copy(root.right);
return x;
}
public LLRBMap<Key, Value> copy() {
Node x = copy(root);
LLRBMap<Key, Value> map = new LLRBMap<Key, Value>();
map.root = x;
return map;
}
}