Java 实现红黑树

public class RedBlackTree<K extends Comparable<K>, V> {

    private static final boolean RED = false;
    private static final boolean BLACK = true;

    private TreeNode<K, V> root;

    public TreeNode<K, V> getRoot() {
        return root;
    }

    public void setRoot(TreeNode<K, V> root) {
        this.root = root;
    }

    static class TreeNode<K, V> {
        private TreeNode<K, V> parent;
        private TreeNode<K, V> left;
        private TreeNode<K, V> right;

        private boolean color = BLACK;
        private K key;
        private V value;

        public TreeNode(TreeNode<K, V> parent, K key, V value) {
            this.parent = parent;
            this.key = key;
            this.value = value;
        }

        public TreeNode(TreeNode<K, V> parent, TreeNode<K, V> left, TreeNode<K, V> right, boolean color, K key, V value) {
            this.parent = parent;
            this.left = left;
            this.right = right;
            this.color = color;
            this.key = key;
            this.value = value;
        }
    }

    /**
     * 新增结点
     * 1. 普通二叉树的插入
     * ① 先查找插入的位置
     * ② 将 K，V 封装成TreeNode对象，插入到红黑树中
     * 2. 红黑树的平衡（旋转+变色）
     */
    public void put(K key, V value) {
        if (key == null) {
            throw new NullPointerException();
        }
        TreeNode<K, V> t = this.root;
        if (t == null) {
            // 红黑树为空，直接创建结点作为根节点
            this.root = new TreeNode<>(null, key, value);
            setColor(root, BLACK);
            return;
        }
        // 查找插入的位置
        int cmp;
        // 记录寻找结点的父节点
        TreeNode<K, V> parent;

        do {
            parent = t;
            cmp = key.compareTo(t.key);
            if (cmp < 0) {
                t = t.left;
            } else if (cmp > 0) {
                t = t.right;
            } else {
                t.value = value;
                return;
            }
        } while (t != null);

        TreeNode<K, V> node = new TreeNode<>(parent, key, value);
        if (cmp < 0) {
            parent.left = node;
        } else {
            parent.right = node;
        }

        // 平衡
        balanceInsertion(node);

    }

    public V remove(K key) {
        TreeNode<K, V> removed = find(key);
        if (removed == null) {
            return null;
        }
        V oldValue = removed.value;
        delete(removed);
        return oldValue;
    }

    private void delete(TreeNode<K, V> removed) {
        if (leftOf(removed) != null && rightOf(removed) != null) {
            TreeNode<K, V> successor = successor(removed);
            removed.key = successor.key;
            removed.value = successor.value;
            removed = successor;
        }
        TreeNode<K, V> replacement = removed.left != null ? removed.left : removed.right;
        if (replacement != null) {
            replacement.parent = removed.parent;
            if (removed.parent == null) {
                root = replacement;
            } else if (removed == leftOf(parentOf(removed))) {
                parentOf(removed).left = replacement;
            } else {
                parentOf(removed).right = replacement;
            }

            removed.left = removed.right = removed.parent = null;
            if (removed.color == BLACK) {
                balanceDeletion(replacement);
            }
        } else if (removed.parent == null) {
            root = null;
        } else {
            if (removed.color == BLACK) {
                balanceDeletion(removed);
            }
            if (removed == leftOf(parentOf(removed))) {
                parentOf(removed).left = null;
            } else {
                parentOf(removed).right = null;
            }
            removed = null;
        }
    }

    private TreeNode<K, V> find(K key) {
        TreeNode<K, V> r = this.root;
        while (r != null) {
            int cmp = key.compareTo(r.key);
            if (cmp < 0) {
                r = r.left;
            } else if (cmp > 0) {
                r = r.right;
            } else {
                return r;
            }
        }
        return null;
    }

    private TreeNode<K, V> predecessor(TreeNode<K, V> node) {
        if (node == null) {
            return null;
        } else if (node.left != null) {
            TreeNode<K, V> p = node.left;
            while (p.right != null) {
                p = p.right;
            }
            return p;
        } else {
            TreeNode<K, V> parent = node.parent;
            TreeNode<K, V> child = node;
            while (parent != null && child == parent.left) {
                child = parent;
                parent = parent.parent;
            }
            return parent;
        }
    }

    private TreeNode<K, V> successor(TreeNode<K, V> node) {
        if (node == null) {
            return null;
        } else if (node.right != null) {
            TreeNode<K, V> p = node.left;
            while (p.left != null) {
                p = p.left;
            }
            return p;
        } else {
            TreeNode<K, V> parent = node.parent;
            TreeNode<K, V> child = node;
            while (parent != null && child == parent.right) {
                child = parent;
                parent = parent.parent;
            }
            return parent;
        }
    }

    /**
     * 左旋的实现
     * p                   r
     * / \                 / \
     * l   r     ===>      p   rr
     * / \             / \
     * rl   rr          l  rl
     */
    private void rotateLeft(TreeNode<K, V> p) {
        TreeNode<K, V> r, pp, rl;

        if (p != null && (r = p.right) != null) {
            if ((rl = p.right = r.left) != null) {
                rl.parent = p;
            }
            if ((pp = r.parent = p.parent) == null) {
                (root = r).color = BLACK;
            } else if (pp.left == p) {
                pp.left = r;
            } else {
                pp.right = r;
            }
            r.left = p;
            p.parent = r;
        }
    }

    /**
     * 右旋的实现
     * p                     l
     * / \                   / \
     * l   r     ===>       ll   p
     * / \                       / \
     * ll  rr                     lr  r
     */
    private void rotateRight(TreeNode<K, V> p) {
        TreeNode<K, V> l, pp, lr;

        if (p != null && (l = p.left) != null) {
            if ((lr = p.left = l.right) != null) {
                lr.parent = p;
            }
            if ((pp = l.parent = p.parent) == null) {
                (root = l).color = BLACK;
            } else if (pp.right == p) {
                pp.right = l;
            } else {
                pp.left = l;
            }
            l.right = p;
            p.parent = l;
        }
    }

    /**
     * 处理颜色，叶结点是黑色
     */
    private boolean colorOf(TreeNode<K, V> node) {
        return node == null ? BLACK : node.color;
    }

    private TreeNode<K, V> parentOf(TreeNode<K, V> node) {
        return node == null ? null : node.parent;
    }

    private TreeNode<K, V> leftOf(TreeNode<K, V> node) {
        return node == null ? null : node.left;
    }

    private TreeNode<K, V> rightOf(TreeNode<K, V> node) {
        return node == null ? null : node.right;
    }

    private void setColor(TreeNode<K, V> node, boolean color) {
        if (node != null) {
            node.color = color;
        }
    }

    /**
     * 插入结点后平衡
     */
    private void balanceInsertion(TreeNode<K, V> insert) {
        // 插入的结点都是红色结点
        setColor(insert, RED);
        TreeNode<K, V> parent, uncle, grandparent;
        while (insert != null && insert != root && parentOf(insert).color == RED) {
            parent = parentOf(insert);
            grandparent = parentOf(parent);

            // insert 的父结点是爷爷结点的左节点
            if (parent == leftOf(grandparent)) {
                uncle = rightOf(grandparent);
                if (colorOf(uncle) == RED) {
                    setColor(parent, BLACK);
                    setColor(uncle, BLACK);
                    setColor(grandparent, RED);
                    insert = grandparent;
                } else {
                    if (insert == rightOf(parent)) {
                        insert = parent;
                        rotateLeft(insert);
                    }
                    setColor(parent, BLACK);
                    setColor(grandparent, RED);
                    rotateRight(grandparent);
                }
            }
            // insert 的父结点是爷爷结点的右节点
            else {
                uncle = leftOf(grandparent);
                if (colorOf(uncle) == RED) {
                    setColor(parent, BLACK);
                    setColor(uncle, BLACK);
                    setColor(grandparent, RED);
                    insert = grandparent;
                } else {
                    if (insert == leftOf(parent)) {
                        insert = parent;
                        rotateRight(insert);
                    }
                    setColor(parent, BLACK);
                    setColor(grandparent, RED);
                    rotateLeft(grandparent);
                }

            }
        }

        // 根节点都为黑色结点
        setColor(root, BLACK);
    }

    private void balanceDeletion(TreeNode<K, V> removed) {
        if (removed != root && colorOf(removed) == BLACK) {
            if (removed == leftOf(parentOf(removed))) {
                TreeNode<K, V> sibling = rightOf(parentOf(removed));
                if (colorOf(sibling) == RED) {
                    setColor(sibling, BLACK);
                    setColor(parentOf(removed), RED);
                    rotateLeft(parentOf(removed));
                    sibling = rightOf(parentOf(removed));
                }
                if (colorOf(leftOf(sibling)) == BLACK && colorOf(parentOf(sibling)) == BLACK) {
                    setColor(sibling, RED);
                    removed = parentOf(removed);
                } else {
                    if (colorOf(rightOf(sibling)) == BLACK) {
                        setColor(sibling, RED);
                        setColor(leftOf(sibling), BLACK);
                        rotateRight(sibling);
                        sibling = rightOf(parentOf(removed));
                    }
                    setColor(sibling, colorOf(parentOf(removed)));
                    setColor(parentOf(removed), BLACK);
                    setColor(rightOf(sibling), BLACK);
                    rotateLeft(parentOf(removed));
                    removed = root;
                }
            } else {
                TreeNode<K, V> sibling = leftOf(parentOf(removed));
                if (colorOf(sibling) == RED) {
                    setColor(sibling, BLACK);
                    setColor(parentOf(removed), RED);
                    rotateRight(parentOf(removed));
                    sibling = leftOf(parentOf(removed));
                }
                if (colorOf(leftOf(sibling)) == BLACK && colorOf(parentOf(sibling)) == BLACK) {
                    setColor(sibling, RED);
                    removed = parentOf(removed);
                } else {
                    if (colorOf(leftOf(sibling)) == BLACK) {
                        setColor(sibling, RED);
                        setColor(rightOf(sibling), BLACK);
                        rotateLeft(sibling);
                        sibling = leftOf(parentOf(removed));
                    }
                    setColor(sibling, colorOf(parentOf(removed)));
                    setColor(parentOf(removed), BLACK);
                    setColor(leftOf(sibling), BLACK);
                    rotateRight(parentOf(removed));
                    removed = root;
                }
            }
        }
        setColor(removed, BLACK);
    }
}