算法导论读书笔记（15）

算法导论读书笔记（15） - 红黑树的具体实现

目录

• 红黑树的简单Java实现

红黑树的简单Java实现

/**
 * 红黑树
 *
 * 部分代码参考自TreeMap源码
 */
public class RedBlackTree<T> {
    protected TreeNode<T> root = null;
    private final Comparator<? super T> comparator;
    private int size = 0;
    private static final boolean RED = false;
    private static final boolean BLACK = true;

    public RedBlackTree() {
        this.comparator = null;
    }

    public RedBlackTree(Comparator<? super T> comparator) {
        this.comparator = comparator;
    }

    public void insert(T key) {
        TreeNode<T> x = root;
        TreeNode<T> y = null;
        TreeNode<T> z = new TreeNode<T>(key, null);
        int cmp;
        if (x == null) {
            root = z;
            size = 1;
            return;
        }
        while (x != null) {
            y = x;
            cmp = compareKey(z.key, x.key);
            if (cmp < 0)
                x = x.left;
            else
                x = x.right;
        }
        z.parent = y;
        cmp = compareKey(z.key, y.key);
        if (cmp < 0)
            y.left = z;
        else
            y.right = z;
        z.left = z.right = null;
        fixAfterInsertion(z);
        size++;
    }

    public T remove(T key) {
        TreeNode<T> p = find(key);
        if (p == null)
            return null;
        T oldValue = p.key;
        deleteNode(p);
        return oldValue;
    }

    public boolean isEmpty() {
        return size() == 0;
    }

    public int size() {
        return size;
    }

    public TreeNode<T> firstNode() {
        return getFirstNode(root);
    }

    public TreeNode<T> lastNode() {
        return getLastNode(root);
    }

    public TreeNode<T> find(T t) {
        TreeNode<T> p = root;
        while (p != null) {
            int cmp = compareKey(t, p.key);
            if (cmp < 0)
                p = p.left;
            else if (cmp > 0)
                p = p.right;
            else
                return p;
        }
        return null;
    }

    public Set<TreeNode<T>> nodeSet() {
        return new NodeSet();
    }

    public static final class TreeNode<T> {
        T key;
        TreeNode<T> left = null;
        TreeNode<T> right = null;
        TreeNode<T> parent;
        boolean color = BLACK;

        TreeNode(T key, TreeNode<T> parent) {
            this.key = key;
            this.parent = parent;
        }

        public T getKey() {
            return key;
        }

        public boolean equals(Object o) {
            if (!(o instanceof TreeNode))
                return false;
            TreeNode<?> e = (TreeNode<?>) o;
            return keyEquals(key, e.getKey());
        }

        public String toString() {
            return "[" + key + " : " + (color ? "BLACK" : "RED") + "]";
        }
    }

    private static <T> TreeNode<T> successor(TreeNode<T> t) {
        if (t == null)
            return t;
        else if (t.right != null)
            return getFirstNode(t.right);
        else {
            TreeNode<T> p = t.parent;
            TreeNode<T> ch = t;
            while (p != null && ch == p.right) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }

    private static <T> TreeNode<T> predecessor(TreeNode<T> t) {
        if (t == null)
            return t;
        else if (t.left != null)
            return getLastNode(t.left);
        else {
            TreeNode<T> p = t.parent;
            TreeNode<T> ch = t;
            while (p != null && ch == p.left) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }

    private static <T> TreeNode<T> getFirstNode(TreeNode<T> t) {
        TreeNode<T> p = t;
        if (p != null)
            while (p.left != null)
                p = p.left;
        return p;
    }

    private static <T> TreeNode<T> getLastNode(TreeNode<T> t) {
        TreeNode<T> p = t;
        if (p != null)
            while (p.right != null)
                p = p.right;
        return p;
    }

    private static boolean keyEquals(Object o1, Object o2) {
        return (o1 == null ? o2 == null : o1.equals(o2));
    }

    private int compareKey(T key1, T key2) {
        int cmp;
        if (comparator != null)
            cmp = comparator.compare(key1, key2);
        else {
            if (key1 == null || key2 == null)
                throw new NullPointerException();
            Comparable<? super T> k = (Comparable<? super T>) key1;
            cmp = k.compareTo(key2);
        }
        return cmp;
    }

    private void rotateLeft(TreeNode<T> p) {
        if (p != null) {
            TreeNode<T> r = p.right;
            p.right = r.left;
            if (r.left != null)
                r.left.parent = p;
            r.parent = p.parent;
            if (p.parent == null)
                root = r;
            else if (p.parent.left == p)
                p.parent.left = r;
            else
                p.parent.right = r;
            r.left = p;
            p.parent = r;
        }
    }

    private void rotateRight(TreeNode<T> p) {
        if (p != null) {
            TreeNode<T> l = p.left;
            p.left = l.right;
            if (l.right != null)
                l.right.parent = p;
            l.parent = p.parent;
            if (p.parent == null)
                root = l;
            else if (p.parent.right == p)
                p.parent.right = l;
            else
                p.parent.left = l;
            l.right = p;
            p.parent = l;
        }
    }

    private void fixAfterInsertion(TreeNode<T> x) {
        x.color = RED;
        while (x != null && x != root && x.parent.color == RED) {
            if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
                TreeNode<T> y = rightOf(parentOf(parentOf(x)));
                if (colorOf(y) == RED) {
                    setColor(parentOf(x), BLACK);
                    setColor(y, BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    x = parentOf(parentOf(x));
                } else {
                    if (x == rightOf(parentOf(x))) {
                        x = parentOf(x);
                        rotateLeft(x);
                    }
                    setColor(parentOf(x), BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    rotateRight(parentOf(parentOf(x)));
                }
            } else {
                TreeNode<T> y = leftOf(parentOf(parentOf(x)));
                if (colorOf(y) == RED) {
                    setColor(parentOf(x), BLACK);
                    setColor(y, BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    x = parentOf(parentOf(x));
                } else {
                    if (x == leftOf(parentOf(x))) {
                        x = parentOf(x);
                        rotateRight(x);
                    }
                    setColor(parentOf(x), BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    rotateLeft(parentOf(parentOf(x)));
                }
            }
        }
        root.color = BLACK;
    }

    private void deleteNode(TreeNode<T> p) {
        size--;
        TreeNode<T> y = p;
        TreeNode<T> x;
        boolean y_original_color = colorOf(y);
        if (leftOf(p) == null) {
            x = rightOf(p);
            transplant(p, rightOf(p));
        } else if (rightOf(p) == null) {
            x = leftOf(p);
            transplant(p, leftOf(p));
        } else {
            y = getFirstNode(rightOf(p));
            y_original_color = colorOf(y);
            x = rightOf(y);
            if (parentOf(y) == p)
                x.parent = y;
            else {
                transplant(y, rightOf(y));
                y.right = rightOf(p);
                y.right.parent = y;
            }
            transplant(p, y);
            y.left = leftOf(p);
            y.left.parent = y;
            y.color = colorOf(p);
        }
        if (y_original_color == BLACK)
            fixAfterDeletion(x);
    }

    private void fixAfterDeletion(TreeNode<T> x) {
        while (x != root && colorOf(x) == BLACK) {
            if (x == leftOf(parentOf(x))) {
                TreeNode<T> sib = rightOf(parentOf(x));

                if (colorOf(sib) == RED) {
                    setColor(sib, BLACK);
                    setColor(parentOf(x), RED);
                    rotateLeft(parentOf(x));
                    sib = rightOf(parentOf(x));
                }

                if (colorOf(leftOf(sib)) == BLACK
                    && colorOf(rightOf(sib)) == BLACK) {
                    setColor(sib, RED);
                    x = parentOf(x);
                } else {
                    if (colorOf(rightOf(sib)) == BLACK) {
                        setColor(leftOf(sib), BLACK);
                        setColor(sib, RED);
                        rotateRight(sib);
                        sib = rightOf(parentOf(x));
                    }
                    setColor(sib, colorOf(parentOf(x)));
                    setColor(parentOf(x), BLACK);
                    setColor(rightOf(sib), BLACK);
                    rotateLeft(parentOf(x));
                    x = root;
                }
            } else { // symmetric
                TreeNode<T> sib = leftOf(parentOf(x));

                if (colorOf(sib) == RED) {
                    setColor(sib, BLACK);
                    setColor(parentOf(x), RED);
                    rotateRight(parentOf(x));
                    sib = leftOf(parentOf(x));
                }

                if (colorOf(rightOf(sib)) == BLACK
                    && colorOf(leftOf(sib)) == BLACK) {
                    setColor(sib, RED);
                    x = parentOf(x);
                } else {
                    if (colorOf(leftOf(sib)) == BLACK) {
                        setColor(rightOf(sib), BLACK);
                        setColor(sib, RED);
                        rotateLeft(sib);
                        sib = leftOf(parentOf(x));
                    }
                    setColor(sib, colorOf(parentOf(x)));
                    setColor(parentOf(x), BLACK);
                    setColor(leftOf(sib), BLACK);
                    rotateRight(parentOf(x));
                    x = root;
                }
            }
        }

        setColor(x, BLACK);
    }

    private void transplant(TreeNode<T> u, TreeNode<T> v) {
        if (parentOf(u) == null)
            root = v;
        else if (u == leftOf(parentOf(u)))
            u.parent.left = v;
        else
            u.parent.right = v;
        if (v != null)
            v.parent = u.parent;
    }

    /**
     * 树的平衡操作
     * 
     * 树的实现没有使用哨兵元素，而是使用下列方法处理null的情况
     */
    private static <T> boolean colorOf(TreeNode<T> p) {
        return (p == null ? BLACK : p.color);
    }

    private static <T> TreeNode<T> parentOf(TreeNode<T> p) {
        return (p == null ? null : p.parent);
    }

    private static <T> void setColor(TreeNode<T> p, boolean c) {
        if (p != null)
            p.color = c;
    }

    private static <T> TreeNode<T> leftOf(TreeNode<T> p) {
        return (p == null) ? null : p.left;
    }

    private static <T> TreeNode<T> rightOf(TreeNode<T> p) {
        return (p == null) ? null : p.right;
    }

    final class NodeSet extends AbstractSet<TreeNode<T>> {

        public Iterator<TreeNode<T>> iterator() {
            return new NodeIterator(firstNode());
        }

        public int size() {
            return RedBlackTree.this.size();
        }
    }

    /**
     * 红黑树的迭代器
     */
    final class NodeIterator extends PrivateNodeIterator<TreeNode<T>> {
        NodeIterator(TreeNode<T> first) {
            super(first);
        }

        public TreeNode<T> next() {
            return nextNode();
        }

        public void remove() {
        }
    }

    abstract class PrivateNodeIterator<E> implements Iterator<E> {
        TreeNode<T> next;

        PrivateNodeIterator(TreeNode<T> first) {
            next = first;
        }

        public boolean hasNext() {
            return next != null;
        }

        final TreeNode<T> nextNode() {
            TreeNode<T> e = next;
            if (e == null)
                throw new NoSuchElementException();
            next = successor(e);
            return e;
        }
    }
}