基于LLRB 红黑树的 TreeMap 实现

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