算法导论示例-OrderStatisticTree

/**
 * Introduction to Algorithms, Second Edition 
 * 14.1 Order-Statistic Tree 
 * 
 * 红黑树的条件：
 * 1.每个节点标记为“红”或“黑”。
 * 2.根标记为“黑”。
 * 3.所有叶节点(nil)标记为“黑”。
 * 4.如果一个节点为“红”，则它的两个节点都为“黑”。
 * 5.对每个的节点，从该节点至后继叶节点包含相同数量的“黑”节点。
 * @author 土豆爸爸
 * 
 */
import java.util.ArrayList;
import java.util.List;

public class OrderStatisticTree {
    /**
     * 数据节点
     */
    public static class Node {
        int key;
        Node parent; //父节点
        Node left; //左子节点
        Node right; //右子节点
        Color color; //节点颜色
        int size; //所在子树节点的数量
        public Node(int key) {
            this.key = key;
        }
        
        public String toString() {
            return String.valueOf(key);
        }
    }
    
    /**
     * 颜色
     */
    private enum Color { RED, BLACK};
    
    Node root; //根节点
    
    Node nil; //空节点
    
    
    /**
     * 构造函数
     */
    public OrderStatisticTree() {
        nil = new Node(-1);
        nil.color = Color.BLACK;
        root = nil;       
    }
    
    /**
     * 从树中选择第i个小键值的节点
     * @param i 序号
     * @return 第i个小键值的节点
     */
    public Node select(int i) {
        return select(root, i); //从根节点开始查询
    }
    
    /**
     * 从以x为根的子树中选择第i个小键值的节点。
     * @param x 子树的根
     * @param i 序号
     * @return 第i个小键值的节点
     */
    private Node select(Node x, int i) {
        int r = x.left.size + 1; //当前节点在以x为根的子树中的序号是左子树的节点个数加1
        if(i == r) { //如果相同，则是当前节点
            return x;
        } else if(i < r) { //如果待查找序号比当前节点序号小，在当前节点的左子树中继续查找
            return select(x.left, i);
        } else { //如果待查找序号比当前节点序号大，在当前节点的左子树中继续查找，序号需要调整
            return select(x.right, i - r);
        }
    }
    
    /**
     * 计算节点x的序号。
     * @param x 待查找节点
     * @return 节点x在树中从小到大排序的序号。
     */
    public int rank(Node x) {
        int r = x.left.size + 1; //当前节点在以x为根的子树中的序号是左子树的节点个数加1
        Node y = x;
        while(y != root) {
            if(y == y.parent.right) { //如果y是右子节点
                r = r + y.parent.left.size + 1; //序号需要加上左兄弟子树的数量，再加父节点
            }
            y = y.parent;
        }
        return r;
    }
    
    /**
     * 左旋转。
     * @param x 支点
     */
    private void leftRotate(Node x) {
        Node y = x.right; // y是x的右子节点
        x.right = y.left; // y的左子树转换成x的右子树
        if (y.left != nil)
            y.left.parent = x;
        y.parent = x.parent; // 用y替换x的位置
        if (x.parent == nil) {
            root = y;
        } else if (x == x.parent.left) {
            x.parent.left = y;
        } else {
            x.parent.right = y;
        }

        y.left = x; // 将x设置为y的左子节点
        x.parent = y;
        y.size = x.size; //y替代x的位置
        x.size = x.left.size + x.right.size + 1; //重新计算x
    }
    
    /**
     * 右旋转。
     * @param x 支点
     */
    private void rightRotate(Node y) {
        Node x = y.left; // x是y的右子节点
        y.left = x.right; // x的右子树转换成y的左子树
        if (x.right != nil)
            x.right.parent = y;
        x.parent = y.parent; // 用x替换y的位置
        if (y.parent == nil) {
            root = x;
        } else if (y == y.parent.left) {
            y.parent.left = x;
        } else {
            y.parent.right = x;
        }

        x.right = y; // 将y设置为x的右子节点
        y.parent = x;
        x.size = y.size; //x替换y的位置
        y.size = y.left.size + y.right.size + 1; //重新计算y
    }
  
    /**
     * 采用递归法查找键值为k的节点。
     * @param k 节点的键值
     * @return 返回键值为k的节点
     */
    public Node search(int k) {
        return search(root, k);
    }
    
    /**
     * 采用递归法查找键值为k的节点。
     * @param x 当前节点
     * @param k 节点的键值
     * @return 返回键值为k的节点
     */
    private Node search(Node x, int k) {
        if(x == nil || k == x.key) {
            return x;
        } else if(k < x.key) {
            return search(x.left, k);
        } else {
            return search(x.right, k);
        }
    }
    
    /**
     * 采用迭代法查找键值为k的节点。
     * @param x 当前节点
     * @param k 节点的键值
     * @return 返回键值为k的节点
     */
    public Node iterativeSearch(int k) {
        return iterativeSearch(root, k);
    }
    
    /**
     * 采用迭代法查找键值为k的节点。
     * @param x 当前节点
     * @param k 节点的键值
     * @return 返回键值为k的节点
     */
    private Node iterativeSearch(Node x, int k) {
        while(x != nil && k != x.key) {
            if(k < x.key) {
                x = x.left;
            } else {
                x = x.right;
            }
        }
        return x;
    }
    
    /**
     * 返回树的最小键值的节点。
     * @return 最小键值的节点
     */
    public Node minimum() {
        return minimum(root);
    }
    
    /**
     * 返回树的最小键值的节点。
     * @param x 当前节点
     * @return 最小键值的节点
     */
    private Node minimum(Node x) {
        while(x.left != nil) {
            x = x.left;
        }
        return x;
    }
    
    /**
     * 返回树的最大键值的节点。
     * @return 最大键值的节点
     */
    public Node maximum() {
        return maximum(root);
    }
    
    /**
     * 返回树的最大键值的节点。
     * @param x 当前节点
     * @return 最大键值的节点
     */
    private Node maximum(Node x) {
        while(x.right != nil) {
            x = x.right;
        }
        return x;
    }
    
    /**
     * 返回指定节点x的后继节点。
     * @param x 当前节点
     * @return x的后继节点；如果x具有最大键值，返回null
     */
    public Node successor(Node x) {
        if(x.right != nil) {
            return minimum(x.right);
        }
        Node y = x.parent;
        while(y != nil && x == y.right) {
            x = y;
            y = y.parent;
        }
        return y;
    }
    
    /**
     * 返回指定节点x的前驱节点。
     * @param x 当前节点
     * @return x的前驱节点；如果x具有最小键值，返回null
     */
    public Node predecessor(Node x) {
        if(x.left != nil) {
            return maximum(x.left);
        }
        Node y = x.parent;
        while(y != nil && x == y.left) {
            x = y;
            y = y.parent;
        }
        return y;
    }
    
    /**
     * 插入节点。
     * @param z 待插入节点
     */
    public void insert(Node z) {
        Node y = nil; //当前节点的父节点
        Node x = root; //当前节点
        while(x != nil) { //迭代查寻z应该所在的位置
            y = x;
            y.size++; //沿着查找路径，将z的所有先辈的size加1
            if(z.key < x.key) {
                x = x.left;
            } else {
                x = x.right;
            }
        }
        z.parent = y;
        z.size = 1; //z是叶节点
        if(y == nil) { 
            root = z; //如果没有父节点，则插入的节点是根节点。
        } else if(z.key < y.key) {
            y.left = z;
        } else {
            y.right = z;
        }
        z.left = nil;
        z.right = nil;
        z.color = Color.RED;
        insertFixup(z);
    }
    
    /**
     * 按红黑树规则进行调整。
     * @param z 待插入节点
     */
    public void insertFixup(Node z) {
        while(z.parent.color == Color.RED) { //违反条件4，并且保证z有爷爷
            if(z.parent == z.parent.parent.left) { //z的父节点是左子节点
                Node y = z.parent.parent.right;
                if(y.color == Color.RED) { //如果z的叔叔是红
                    z.parent.color = Color.BLACK; //将z的父亲和叔叔设为黑
                    y.color = Color.BLACK;
                    z.parent.parent.color = Color.RED; //z的爷爷设为红
                    z = z.parent.parent; //迭代
                } else { //如果z的叔叔是黑
                    if (z == z.parent.right) { //如果z是右子节点，左旋
                        z = z.parent;
                        leftRotate(z);
                    }
                    z.parent.color = Color.BLACK; //z的父亲为黑(叔叔为黑)
                    z.parent.parent.color = Color.RED; //z的爷爷为红
                    rightRotate(z.parent.parent); // 右旋
                }
            } else { //z的父节点是右子节点，反向对称
                Node y = z.parent.parent.left;
                if(y.color == Color.RED) {
                    z.parent.color = Color.BLACK;
                    y.color = Color.BLACK;
                    z.parent.parent.color = Color.RED;
                    z = z.parent.parent;
                } else {
                    if (z == z.parent.left) {
                        z = z.parent;
                        rightRotate(z);
                    }
                    z.parent.color = Color.BLACK;
                    z.parent.parent.color = Color.RED;
                    leftRotate(z.parent.parent);
                }
            }
        }
        root.color = Color.BLACK; //满足条件2
    }
    
    /**
     * 删除节点。
     * @param z 待删除节点
     */
    public Node delete(Node z) {
        Node y = null;
        Node x = null;
        if (z.left == nil || z.right == nil) {
            y = z;
        } else {
            y = successor(z);
        }

        if (y.left != nil) {
            x = y.left;
        } else {
            x = y.right;
        }
        
        x.parent = y.parent;
        
        if (y.parent == nil) {
            root = x;
        } else if (y == y.parent.left) {
            y.parent.left = x;
        } else {
            y.parent.right = x;
        }

        Node p = y.parent; //调整y所有父节点的size
        while(p != nil) {
            p.size--;
            p = p.parent;
        }
        
        if (y != z) { // 如果z包含两个子节点，用y替换z的位置
            y.parent = z.parent;
            if (z.parent != nil) {
                if (z == z.parent.left) {
                    z.parent.left = y;
                } else {
                    z.parent.right = y;
                }
            } else {
                root = y;
            }
            z.left.parent = y;
            y.left = z.left;
            z.right.parent = y;
            y.right = z.right;
            y.size = y.left.size + y.right.size + 1; //重新计算y的size
        }
        
        if(y.color == Color.BLACK) {
            deleteFixup(x);
        }
        return y;
    }
    
    /**
     * 按红黑树规则进行调整。
     * @param z 待删除节点
     */
    private void deleteFixup(Node x) {
        while (x != nil && x != root && x.color == Color.BLACK) {
            if (x == x.parent.left) {
                Node w = x.parent.right;
                if (w == nil)
                    return;
                if (w.color == Color.RED) {
                    w.color = Color.BLACK;
                    x.parent.color = Color.RED;
                    leftRotate(x.parent);
                    w = x.parent.right;
                }
                if (w == nil)
                    return;
                if (w.left.color == Color.BLACK && w.right.color == Color.BLACK) {
                    w.color = Color.RED;
                    x = x.parent;
                } else {
                    if (w.right.color == Color.BLACK) {
                        w.left.color = Color.BLACK;
                        w.color = Color.RED;
                        rightRotate(w);
                        w = x.parent.right;
                    }

                    w.color = x.parent.color;
                    x.parent.color = Color.BLACK;
                    w.right.color = Color.BLACK;
                    leftRotate(x.parent);
                    x = root;
                }
            } else {
                Node w = x.parent.left;
                if (w == nil)
                    return;
                if (w.color == Color.RED) {
                    w.color = Color.BLACK;
                    x.parent.color = Color.RED;
                    rightRotate(x.parent);
                    w = x.parent.left;
                }
                if (w == nil)
                    return;
                if (w.right.color == Color.BLACK && w.left.color == Color.BLACK) {
                    w.color = Color.RED;
                    x = x.parent;
                } else {
                    if (w.left.color == Color.BLACK) {
                        w.right.color = Color.BLACK;
                        w.color = Color.RED;
                        leftRotate(w);
                        w = x.parent.left;
                    }
                    w.color = x.parent.color;
                    x.parent.color = Color.BLACK;
                    w.left.color = Color.BLACK;
                    rightRotate(x.parent);
                    x = root;

                }
            }
        }
        x.color = Color.BLACK;
    }
    
    /**
     * 中序遍历。即从小到大排序。
     * @return 返回已排序的节点列表
     */
    public List<Node> inorderWalk() {
        List<Node> list = new ArrayList<Node>();
        inorderWalk(root, list);
        return list;
    }
    
    /**
     * 中序遍历。
     * @param x 当前节点
     * @param list 遍历结果存储在list中
     */
    private void inorderWalk(Node x, List<Node> list) {
        if(x != nil) {
            inorderWalk(x.left, list);
            list.add(x);
            inorderWalk(x.right, list);
        }
    }
    
    /**
     * 前序遍历打印。即从小到大排序。
     * @return 返回已排序的节点列表
     */
    public void preorderWalk() {
        preorderWalk(root);
    }
    
    /**
     * 中序遍历打印。
     * @param x 当前节点
     * @param list 遍历结果存储在list中
     */
    private void preorderWalk(Node x) {
        if(x != nil) {
            System.out.print("(");
            System.out.print(x);
            preorderWalk(x.left);
            preorderWalk(x.right);
            System.out.print(")");
        } else {
            System.out.print("N");
        }
    }
}

import java.util.List;

import junit.framework.TestCase;

public class OrderStatisticTreeTest extends TestCase {
    public void testLinkedList() {
        //while (true) {
            OrderStatisticTree tree = new OrderStatisticTree();
            // 插入N个随机节点
            int count = 30;
            for (int i = 0; i < count;) {
                int key = (int) (Math.random() * 100);
                if (tree.search(key) == tree.nil) {
                    tree.insert(new OrderStatisticTree.Node(key));                    
                    i++;
                    assertEquals(i, tree.root.size);
                }
            }
            // 测试最大值,最小值
            List<OrderStatisticTree.Node> list = tree.inorderWalk();
            verifyOrdered(list);
            assertEquals(count, list.size());
            assertEquals(list.get(0), tree.minimum());
            assertEquals(list.get(list.size() - 1), tree.maximum());

            // 测试后继
            OrderStatisticTree.Node min = tree.minimum(), succ = null;
            while ((succ = tree.successor(min)) != tree.nil) {
                assertTrue(succ.key > min.key);
                min = succ;
            }

            // 测试前驱
            OrderStatisticTree.Node max = tree.maximum(), pre = null;
            while ((pre = tree.predecessor(max)) != tree.nil) {
                assertTrue(succ.key < max.key);
                max = pre;
            }

            // 测试删除
            int[] keys = new int[list.size()];
            for (int i = 0; i < keys.length; i++) {
                keys[i] = list.get(i).key;
                //测试rank
                assertEquals(i + 1, tree.rank(tree.search(keys[i])));
                //测试select
                assertEquals(tree.search(keys[i]), tree.select(i + 1));
            }
            
            PermuteBySorting.permute(keys);
            for (int i = 0; i < count; i++) {
                OrderStatisticTree.Node node = tree.search(keys[i]);
                assertNotNull(node);
                tree.delete(node);
                assertEquals(tree.nil, tree.search(keys[i]));
                verifyOrdered(tree.inorderWalk());
            }
        //}
    }

    private boolean verifyOrdered(List<OrderStatisticTree.Node> list) {
        for (int i = 1; i < list.size(); i++) {
            if (list.get(i - 1).key > list.get(i).key) {
                return false;
            }
        }
        return true;
    }
}