《算法》第三章部分程序 part 3

▶ 书中第三章部分程序，加上自己补充的代码，红黑树

● 红黑树，大部分方法与注释与二叉树相同

package package01;

import java.util.NoSuchElementException;
import edu.princeton.cs.algs4.Queue;
import edu.princeton.cs.algs4.StdIn;
import edu.princeton.cs.algs4.StdOut;

public class class01<Key extends Comparable<Key>, Value>
{
    private static final boolean RED = true;
    private static final boolean BLACK = false;

    private class Node
    {
        private Key key;
        private Value val;
        private Node left, right;
        private boolean color;     // 指向该节点的连接的颜色
        private int size;

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

    private Node root;

    public class01() {}

    private boolean isRed(Node x)
    {
        if (x == null)
            return false;
        return x.color == RED;
    }

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

    private int size(Node x)
    {
        if (x == null)
            return 0;
        return x.size;
    }

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

    public Value get(Key key)               // 查找
    {
        if (key == null)
            throw new IllegalArgumentException("
<get> key == null.
");
        return getKernel(root, key);
    }

    private Value getKernel(Node x, Key key)// 查找内核，使用了非递归实现
    {
        for (int cmp = key.compareTo(x.key); x != null; cmp = key.compareTo(x.key))
        {
            if (cmp < 0)
                x = x.left;
            else if (cmp > 0)
                x = x.right;
            else
                return x.val;
        }
        return null;
    }

    public boolean contains(Key key)        // 判断 key 是否在树中
    {
        if (key == null)
            throw new IllegalArgumentException("
<contains> key == null.
");
        return get(key) != null;
    }

    public void put(Key key, Value val)     // 插入
    {
        if (key == null)
            throw new IllegalArgumentException("
<put> key == null.
");
        if (val == null)
            delete(key);
        else
        {
            root = putKernel(root, key, val);
            root.color = BLACK;
        }
        // assert check();
    }

    private Node putKernel(Node x, Key key, Value val)
    {
        if (x == null)
            return new Node(key, val, RED, 1);
        int cmp = key.compareTo(x.key);
        if (cmp < 0)
            x.left = putKernel(x.left, key, val);
        else if (cmp > 0)
            x.right = putKernel(x.right, key, val);
        else
            x.val = val;
        if (isRed(x.right) && !isRed(x.left))
            x = rotateLeft(x);
        if (isRed(x.left) && isRed(x.left.left))
            x = rotateRight(x);
        if (isRed(x.left) && isRed(x.right))
            flipColors(x);
        x.size = 1 + size(x.left) + size(x.right);
        return x;
    }

    public void deleteMin()
    {
        if (isEmpty())
            throw new NoSuchElementException("
<deleteMin> underflow.
");
        if (!isRed(root.left) && !isRed(root.right))
            root.color = RED;
        root = deleteMinKernel(root);
        if (!isEmpty())
            root.color = BLACK;
        // assert check();
    }

    private Node deleteMinKernel(Node x)
    {
        if (x.left == null)
            return null;
        if (!isRed(x.left) && !isRed(x.left.left))
            x = moveRedLeft(x);
        x.left = deleteMinKernel(x.left);
        return balance(x);
    }

    public void deleteMax()
    {
        if (isEmpty())
            throw new NoSuchElementException("
<deleteMax> underflow.
");
        if (!isRed(root.left) && !isRed(root.right))
            root.color = RED;
        root = deleteMaxKernel(root);
        if (!isEmpty())
            root.color = BLACK;
        // assert check();
    }

    private Node deleteMaxKernel(Node x)
    {
        if (isRed(x.left))
            x = rotateRight(x);
        if (x.right == null)
            return null;
        if (!isRed(x.right) && !isRed(x.right.left))
            x = moveRedRight(x);
        x.right = deleteMaxKernel(x.right);
        return balance(x);
    }

    public void delete(Key key)
    {
        if (key == null)
            throw new IllegalArgumentException("
<delete> key == null.
");
        if (!contains(key))
            return;
        if (!isRed(root.left) && !isRed(root.right))
            root.color = RED;
        root = deleteKernel(root, key);
        if (!isEmpty())
            root.color = BLACK;
        // assert check();
    }

    private Node deleteKernel(Node x, Key key)
    {
        if (key.compareTo(x.key) < 0)
        {
            if (!isRed(x.left) && !isRed(x.left.left))
                x = moveRedLeft(x);
            x.left = deleteKernel(x.left, key);
        }
        else
        {
            if (isRed(x.left))
                x = rotateRight(x);
            if (key.compareTo(x.key) == 0 && (x.right == null))
                return null;
            if (!isRed(x.right) && !isRed(x.right.left))
                x = moveRedRight(x);
            if (key.compareTo(x.key) == 0)
            {
                Node t = minKernel(x.right);
                x.key = t.key;
                x.val = t.val;
                // x.val = get(x.right, min(x.right).key);
                // x.key = min(x.right).key;
                x.right = deleteMinKernel(x.right);
            }
            else x.right = deleteKernel(x.right, key);
        }
        return balance(x);
    }

    private Node rotateRight(Node x)    // 右旋转
    {
        Node t = x.left;
        x.left = t.right;
        t.right = x;
        t.color = t.right.color;
        t.right.color = RED;
        t.size = x.size;
        x.size = 1 + size(x.left) + size(x.right);
        return t;
    }

    private Node rotateLeft(Node x)     // 左旋转
    {
        Node t = x.right;
        x.right = t.left;
        t.left = x;
        t.color = t.left.color;
        t.left.color = RED;
        t.size = x.size;
        x.size = 1 + size(x.left) + size(x.right);
        return t;
    }

    private void flipColors(Node x)     // 改变节点及其子节点的颜色
    {
        // assert (x != null) && (x.left != null) && (x.right != null);
        // assert (!isRed(x) &&  isRed(x.left) &&  isRed(x.right)) || (isRed(x)  && !isRed(x.left) && !isRed(x.right));
        x.color = !x.color;
        x.left.color = !x.left.color;
        x.right.color = !x.right.color;
    }

    private Node moveRedLeft(Node x)    // x 红而 x.left 和 x.left.left 都是黑的，调整使得两个黑链接之一变红
    {
        // assert (x != null);
        // assert isRed(x) && !isRed(x.left) && !isRed(x.left.left);
        flipColors(x);
        if (isRed(x.right.left))
        {
            x.right = rotateRight(x.right);
            x = rotateLeft(x);
            flipColors(x);
        }
        return x;
    }

    private Node moveRedRight(Node x)    // x 红而 x.right 和 x.right.left 都是黑的，调整使得两黑链接之一变红
    {
        // assert (x != null);
        // assert isRed(x) && !isRed(x.right) && !isRed(x.right.left);
        flipColors(x);
        if (isRed(x.left.left))
        {
            x = rotateRight(x);
            flipColors(x);
        }
        return x;
    }

    private Node balance(Node x)        // 右链接红色、连续两层左链接红色以及左右链接都是红色的状况，分别调整
    {
        // assert (x != null);
        if (isRed(x.right))
            x = rotateLeft(x);
        if (isRed(x.left) && isRed(x.left.left))
            x = rotateRight(x);
        if (isRed(x.left) && isRed(x.right))
            flipColors(x);
        x.size = 1 + size(x.left) + size(x.right);
        return x;
    }

    public Key min()
    {
        if (isEmpty())
            throw new NoSuchElementException("
<delete> empty.
");
        return minKernel(root).key;
    }

    private Node minKernel(Node x)
    {
        if (x.left == null)
            return x;
        return minKernel(x.left);
    }

    public Key max()
    {
        if (isEmpty())
            throw new NoSuchElementException("
<max> empty.
");
        return maxKernel(root).key;
    }

    private Node maxKernel(Node x)
    {
        if (x.right == null)
            return x;
        return maxKernel(x.right);
    }

    public Key floor(Key key)
    {
        if (key == null)
            throw new IllegalArgumentException("
<floor> key == null.
");
        if (isEmpty())
            throw new NoSuchElementException("
<floor> empty.
");
        Node x = floorKernel(root, key);
        return (x == null) ? null : x.key;
    }

    private Node floorKernel(Node x, Key key)
    {
        if (x == null)
            return null;
        int cmp = key.compareTo(x.key);
        if (cmp == 0)
            return x;
        if (cmp < 0)
            return floorKernel(x.left, key);
        Node t = floorKernel(x.right, key);
        return (t == null) ? x : t;
    }

    public Key ceiling(Key key)
    {
        if (key == null)
            throw new IllegalArgumentException("
<ceiling> key == null.
");
        if (isEmpty())
            throw new NoSuchElementException("
<ceiling> empty.
");
        Node x = ceilingKernel(root, key);
        return (x == null) ? null : x.key;
    }

    private Node ceilingKernel(Node x, Key key)
    {
        if (x == null)
            return null;
        int cmp = key.compareTo(x.key);
        if (cmp == 0)
            return x;
        if (cmp > 0)
            return ceilingKernel(x.right, key);
        Node t = ceilingKernel(x.left, key);
        return (t == null) ? x : t;
    }

    public Key select(int k)
    {
        if (k < 0 || k >= size())
            throw new IllegalArgumentException("
<select> k < 0 || k >= size().
");
        return selectKernel(root, k).key;
    }

    private Node selectKernel(Node x, int k)
    {
        if (x == null)
            return null;
        int t = size(x.left);
        if (k <t)
            return selectKernel(x.left, k);
        if (k > t)
            return selectKernel(x.right, k - t - 1);
        return x;
    }

    public int rank(Key key)
    {
        if (key == null)
            throw new IllegalArgumentException("
<rank> key == null.
");
        return rankKernel(key, root);
    }

    private int rankKernel(Key key, Node x)
    {
        if (x == null)
            return 0;
        int cmp = key.compareTo(x.key);
        if (cmp < 0)
            return rankKernel(key, x.left);
        if (cmp > 0)
            return 1 + size(x.left) + rankKernel(key, x.right);
        return size(x.left);
    }

    public Iterable<Key> keys()
    {
        if (isEmpty())
            return new Queue<Key>();
        Key lo = min(), hi = max();
        if (lo == null)
            throw new IllegalArgumentException("
<iterable> lo == null.
");
        if (hi == null)
            throw new IllegalArgumentException("
<iterable> hi == null.
");
        Queue<Key> queue = new Queue<Key>();
        // if (isEmpty() || lo.compareTo(hi) > 0) return queue;
        keysKernel(root, queue, lo, hi);
        return queue;
    }

    private void keysKernel(Node x, Queue<Key> queue, Key lo, Key hi)
    {
        if (x == null)
            return;
        int cmplo = lo.compareTo(x.key), cmphi = hi.compareTo(x.key);
        if (cmplo < 0)
            keysKernel(x.left, queue, lo, hi);
        if (cmplo <= 0 && cmphi >= 0)
            queue.enqueue(x.key);
        if (cmphi > 0)
            keysKernel(x.right, queue, lo, hi);
    }

    public Iterable<Key> levelOrder()
    {
        Queue<Key> keys = new Queue<Key>();
        Queue<Node> queue = new Queue<Node>();
        for (queue.enqueue(root); !queue.isEmpty();)
        {
            Node x = queue.dequeue();
            if (x == null)
                continue;
            keys.enqueue(x.key);
            queue.enqueue(x.left);
            queue.enqueue(x.right);
        }
        return keys;
    }

    public int size(Key lo, Key hi)
    {
        if (lo == null)
            throw new IllegalArgumentException("
<size> lo == null.
");
        if (hi == null)
            throw new IllegalArgumentException("
<size> hi == null.
");
        if (lo.compareTo(hi) > 0)
            return 0;
        if (contains(hi))
            return rank(hi) - rank(lo) + 1;
        return rank(hi) - rank(lo);
    }

    public int height()
    {
        return heightKernel(root);
    }

    private int heightKernel(Node x)
    {
        if (x == null)
            return -1;
        return 1 + Math.max(heightKernel(x.left), heightKernel(x.right));
    }    

    private boolean check()
    {
        if (!isBST())
            StdOut.println("
<check> Not in symmetric order.
");
        if (!isSizeConsistent())
            StdOut.println("
<check> Subtree counts not consistent.
");
        if (!isRankConsistent())
            StdOut.println("
<check> Ranks not consistent.
");
        if (!is23())
            StdOut.println("
<check> Not a 2 - 3 tree.
");
        if (!isBalanced())
            StdOut.println("
<check> Not balanced.
");
        return isBST() && isSizeConsistent() && isRankConsistent() && is23() && isBalanced();
    }

    private boolean isBST()
    {
        return isBSTKernel(root, null, null);
    }

    private boolean isBSTKernel(Node x, Key min, Key max)
    {
        if (x == null)
            return true;
        if (min != null && x.key.compareTo(min) <= 0)
            return false;
        if (max != null && x.key.compareTo(max) >= 0)
            return false;
        return isBSTKernel(x.left, min, x.key) && isBSTKernel(x.right, x.key, max);
    }
    
    private boolean isSizeConsistent()
    {
        return isSizeConsistentKernel(root);
    }

    private boolean isSizeConsistentKernel(Node x)
    {
        if (x == null)
            return true;
        if (x.size != 1 + size(x.left) + size(x.right))
            return false;
        return isSizeConsistentKernel(x.left) && isSizeConsistentKernel(x.right);
    }

    private boolean isRankConsistent()
    {
        for (int i = 0; i < size(); i++)
        {
            if (i != rank(select(i)))
                return false;
        }
        for (Key key : keys())
        {
            if (key.compareTo(select(rank(key))) != 0)
                return false;
        }
        return true;
    }

    private boolean is23()
    {
        return is23Kernel(root);
    }

    private boolean is23Kernel(Node x)
    {
        if (x == null)
            return true;
        if (isRed(x.right))
            return false;
        if (x != root && isRed(x) && isRed(x.left))
            return false;
        return is23Kernel(x.left) && is23Kernel(x.right);
    }

    private boolean isBalanced()
    {
        int black = 0;                          // 从 root 到 min 节点的路径长        
        for(Node x = root;x != null; x = x.left)
        {
            if (!isRed(x)) 
                black++;            
        }
        return isBalancedKernel(root, black);
    }

    private boolean isBalancedKernel(Node x, int black)
    {
        if (x == null) 
            return black == 0;
        if (!isRed(x)) 
            black--;
        return isBalancedKernel(x.left, black) && isBalancedKernel(x.right, black);
    }

    public static void main(String[] args)
    {
        class01<String, Integer> st = new class01<String, Integer>();
        for (int i = 0; !StdIn.isEmpty(); i++)
        {
            String key = StdIn.readString();
            st.put(key, i);
        }

        for (String s : st.levelOrder())
            StdOut.println(s + " " + st.get(s));

        StdOut.println();
        for (String s : st.keys())
            StdOut.println(s + " " + st.get(s));
    }
}// 703