紅黑樹java代碼實現

紅黑樹 思想源於：https://www.cnblogs.com/nananana/p/10434549.html
有解釋有圖，很清晰（刪除時需考慮根節點和兄弟節點的子節點是否存在）

package tree;

import jdk.nashorn.internal.ir.CallNode;

import java.util.Random;

/**
 * 紅黑樹
 * 思想源於：https://www.cnblogs.com/nananana/p/10434549.html
 */
public class RBTree {
    private RBTree.Node root = null;

    private enum Color {RED, BLACK}

    private enum Child {LEFT, RIGHT}

    private class Node {
        private Integer key;    //key
        private Object data;    //value
        private Node leftChild;   //左子節點
        private Node rightChild;  //右子節點
        private Node parent;  //父節點
        private Color color;   //紅黑標示

        private Node() {
        }

        Node(Object key, Object data, Color color) {
            this.key = (Integer) key;
            this.data = data;
            this.color = color;
        }

        public Color getColor() {
            return color;
        }

        public void setColor(Color color) {
            this.color = color;
        }

        boolean isRed() {
            if (this.color.equals(Color.RED))
                return true;
            return false;
        }
    }

    /**
     * 插入數據
     *
     * @param value 插入數據
     * @return 數據重複返回false
     */
    boolean insertNode(Integer key, Object value) {
        return insertNode(root, key, value, null, Child.LEFT);
    }

    private boolean insertNode(Node node, Integer key, Object value, Node preNode, Child child) {
        if (node == null) {
            node = new Node(key, value, Color.RED);
            if (preNode == null) {  //父節點爲空，將node設爲根節點
                root = node;
            } else {
                if (child.equals(Child.LEFT)) {
                    preNode.leftChild = node;
                } else {
                    preNode.rightChild = node;
                }
                node.parent = preNode;
            }

            //經過RB_INSERT_FIXUP對紅黑樹的結點進行顏色修改以及旋轉，讓樹仍然是一顆紅黑樹
            RB_INSERT_FIXUP(node);
            return true;
        } else {
            if (key.compareTo(node.key) == 0) {
                //root = node;
                return false;
            } else if (key.compareTo(node.key) < 0) {
                if (!insertNode(node.leftChild, key, value, node, Child.LEFT)) {
                    return false;
                }
            } else {
                if (!insertNode(node.rightChild, key, value, node, Child.RIGHT)) {
                    return false;
                }
            }
            return true;
        }
    }

    /**
     * @param node 插入節點
     */
    private void RB_INSERT_FIXUP(Node node) {
        Node pNode = node.parent;
        if (node == root) { //插入結點爲跟節點，直接改變顏色
            node.setColor(Color.BLACK);
            return;
        }
        if (node == null || pNode.color.equals(Color.BLACK)) {    //插入結點的父結點爲黑結點,因爲插入的結點是紅色的，並不會影響紅黑樹的平衡，直接插入便可，無需作自平衡。
            return;
        } else {    //情景4：插入結點的父結點爲紅結點
            Node graNode = node.parent.parent;
            if (pNode == graNode.leftChild) {  //父節點是祖父節點的左子節點
                if (graNode.rightChild != null && graNode.rightChild.isRed()) { //情景4.1：叔叔結點存在而且爲紅結點
                    /*將P和S設置爲黑色（當前插入結點I）將gra設置爲紅色 把gra設置爲當前插入結點*/
                    pNode.setColor(Color.BLACK);
                    graNode.rightChild.setColor(Color.BLACK);
                    graNode.setColor(Color.RED);
                    RB_INSERT_FIXUP(graNode);
                } else { //情景4.2：叔叔結點不存在或爲黑結點，而且插入結點的父親結點是祖父結點的左子結點
                    if (node == pNode.leftChild) {//情景4.2.1 插入結點是其父結點的左子結點
                        /*將P設爲黑色 將gra設爲紅色 對gra進行右旋*/
                        pNode.setColor(Color.BLACK);
                        graNode.setColor(Color.RED);
                        RRotate(graNode);
                    } else {    //情景4.2.2 插入結點是其父結點的右子結點
                        /*對P進行左旋 把P設置爲插入結點，獲得情景4.2.1 進行情景4.2.1的處理*/
                        LRotate(pNode);
                        RB_INSERT_FIXUP(pNode);
                    }

                }
            } else { //4.3 父節點是祖父節點的右子節點
                if (graNode.leftChild != null && graNode.leftChild.isRed()) { //情景4.3：叔叔結點存在而且爲紅結點+
                    /*將P和S設置爲黑色（當前插入結點I）將gra設置爲紅色 把gra設置爲當前插入結點*/
                    pNode.setColor(Color.BLACK);
                    graNode.leftChild.setColor(Color.BLACK);
                    graNode.setColor(Color.RED);
                    RB_INSERT_FIXUP(graNode);
                } else { //情景4.3.1：叔叔結點不存在或爲黑結點，而且插入結點的父親結點是祖父結點的左子結點
                    if (node == pNode.rightChild) {//情景4.3.1：插入結點是其父結點的右子結點
                        /*將P設爲黑色 將gra設爲紅色 對PP進行左旋*/
                        pNode.setColor(Color.BLACK);
                        graNode.setColor(Color.RED);
                        LRotate(graNode);
                    } else {    //情景4.3.2 插入結點是其父結點的右子結點
                        /*對P進行右旋 把P設置爲插入結點，獲得情景4.3.1 進行情景4.3.1的處理*/
                        RRotate(pNode);
                        RB_INSERT_FIXUP(pNode);
                    }
                }
            }
        }
    }

    /**
     * 右旋
     *
     * @param T
     */
    private void RRotate(Node T) {
        Node lc = T.leftChild;
        T.leftChild = lc.rightChild;
        if (T.leftChild != null) {
            T.leftChild.parent = T;
        }
        lc.rightChild = T;
        returnPNode(T, lc);
    }

    private Node returnPNode(Node T, Node node) {
        if (T == root) {
            root = node;
        } else if (T.parent.leftChild == T) {
            T.parent.leftChild = node;
        } else {
            T.parent.rightChild = node;
        }
        node.parent = T.parent;
        T.parent = node;
        return node;
    }

    /**
     * 左旋
     *
     * @param T
     */
    private void LRotate(Node T) {
        Node rc = T.rightChild;
        T.rightChild = rc.leftChild;
        if (T.rightChild != null) {

            T.rightChild.parent = T;
        }
        rc.leftChild = T;
        returnPNode(T, rc);
    }

    /**
     * 中序
     */
    public void ldrTraversal() {
        ldrTraversal(root);
    }

    /**
     * 中序
     */
    private void ldrTraversal(Node node) {
        if (node != null) {
            ldrTraversal(node.leftChild);
            System.out.print(node.key + ":" + node.color + ";");
            //System.out.print("key:" + node.key + "-value" + node.data + ":" + node.color + ";");
            ldrTraversal(node.rightChild);
        }

    }

    /**
     * 先序
     */
    public void dlrTraversal() {
        dlrTraversal(root);
    }

    /**
     * 先序
     */
    private void dlrTraversal(Node node) {
        if (node != null) {
            System.out.print(node.key + ":" + node.color + ";");
            dlrTraversal(node.leftChild);
            dlrTraversal(node.rightChild);
        }

    }

    /**
     * 後序
     */
    public void lrdTraversal() {
        lrdTraversal(root);
    }

    /**
     * 後序
     */
    private void lrdTraversal(Node node) {
        if (node != null) {
            lrdTraversal(node.leftChild);
            lrdTraversal(node.rightChild);
            System.out.print("key:" + node.key + "-value" + node.data + ":" + node.color + ";");
        }

    }

    /**
     * 搜索
     *
     * @param key 傳入key
     * @return 返回value
     */
    public Object search(Integer key) {
        if (this.root != null) {
            return searchNode(key, root).data;
        }
        return null;
    }

    /**
     * @param key 刪除key對應的node
     */
    public boolean removen(Integer key) {
        if (this.root != null) {
            Node node = searchNode(key, root);
            if (node == null) {
                return false;
            }
            removenNode(node);
            return true;
        }
        return false;
    }

    /**
     * @param node 刪除的節點
     */
    private void removenNode(Node node) {
        if (node == null) {
            return;
        }
        if (node.leftChild == null && node.rightChild == null) {    //情景1：若刪除結點無子結點，直接刪除。
            changeNode(node, null);
        } else if (node.leftChild != null && node.rightChild != null) { //情景3：若刪除結點有兩個子結點，用後繼結點（大於刪除結點的最小結點）替換刪除結點。
            Node rNode = node.rightChild;
            while (rNode.leftChild != null) {   //找到後繼結點
                rNode = rNode.leftChild;
            }
            //  交換位子
            /*if (rNode == node.rightChild) {
                node.rightChild = null;
                rNode.leftChild = node.leftChild;
            } else {
                if (rNode.rightChild != null) {    //後繼節點若是有右節點
                    rNode.parent.leftChild = rNode.rightChild;
                    rNode.rightChild.parent = rNode.parent;
                }
                rNode.leftChild = node.leftChild;
                rNode.rightChild = node.rightChild;
            }*/
            changeNode(node, rNode);    //用後繼節點替換要刪除節點
        } else { //情景2：若刪除結點只有一個子結點，用子結點替換刪除結點。
            if (node.leftChild != null) {
                changeNode(node, node.leftChild);
            } else {
                changeNode(node, node.rightChild);
            }
        }
    }

    /**
     * 兩節點位置交換
     * 交換後刪除替換節點fixupNode
     * @param delNode   要刪除節點
     * @param fixupNode 替換節點
     */
    private void changeNode(Node delNode, Node fixupNode) {
        RB_DELETE_FIXUP(fixupNode);

        if (fixupNode == null) {
            if (delNode.parent.leftChild == delNode) {
                delNode.parent.leftChild = null;
            } else {
                delNode.parent.rightChild = null;
            }
            return;
        }

        Object data = delNode.data;
        Color color = delNode.color;
        Integer key = delNode.key;
        if (delNode == root) {  // 交換時若是刪除節點是根節點，顏色直接改爲黑色
            delNode.setColor(Color.BLACK);
        } else {
            delNode.color = fixupNode.color;
        }
        delNode.key = fixupNode.key;
        delNode.data = fixupNode.data;
        fixupNode.key = key;
        fixupNode.data = data;
        fixupNode.color = color;

        removenNode(fixupNode);
    }

    public Node searchNode(Integer key, Node node) {
        if (node == null)
            return null;
        if (node.key.compareTo(key) == 0) {
            return node;
        } else if (key.compareTo(node.key) < 0) {
            if (node.leftChild != null) {
                return searchNode(key, node.leftChild);
            }
            return null;
        } else {
            if (node.rightChild != null) {
                return searchNode(key, node.rightChild);
            }
            return null;
        }
    }

    private void RB_DELETE_FIXUP(Node fixupNode) {
        if (fixupNode == null || fixupNode.isRed()) {    //情景1：替換結點是紅色結點
            /*顏色變爲刪除結點的顏色*/
            return;
        } else {  //情景2：替換結點是黑結點
            Node bNode = fixupNode.parent.rightChild;
            if (fixupNode == fixupNode.parent.leftChild) { //情景2.1：替換結點是其父結點的左子結點
                //情景2.1.1：替換結點的兄弟結點是紅結點
                if (bNode.isRed()) {
                    bNode.setColor(Color.BLACK);
                    fixupNode.parent.setColor(Color.RED);
                    RRotate(fixupNode.parent);
                    RB_DELETE_FIXUP(fixupNode);
                } else {  //情景2.1.2： 替換結點的兄弟結點是黑結點
                    //情景2.1.2.1：替換結點的兄弟結點的右子結點是紅結點，左子結點任意顏色
                    if (bNode.rightChild != null && bNode.rightChild.isRed()) {
                        /*將S的顏色設爲P的顏色 將P設爲黑色 將SR設爲黑色 對P進行左旋*/
                        bNode.color = fixupNode.parent.color;
                        fixupNode.parent.setColor(Color.BLACK);
                        bNode.rightChild.setColor(Color.RED);
                        LRotate(fixupNode.parent);
                    } else if (bNode.leftChild != null && bNode.leftChild.isRed()) {
                        //情景2.1.2.2：替換結點的兄弟結點的右子結點爲黑結點，左子結點爲紅結點
                        /*將S設爲紅色 將SL設爲黑色 對S進行右旋，獲得情景2.1.2.1 進行情景2.1.2.1的處理*/
                        bNode.setColor(Color.RED);
                        bNode.leftChild.setColor(Color.BLACK);
                        RRotate(bNode);
                        RB_DELETE_FIXUP(fixupNode);
                    } else {//刪除情景2.1.2.3： 替換結點的兄弟結點的子結點都爲黑結點
                        /*將S設爲紅色 把P做爲新的替換結點 從新進行刪除結點情景處理*/
                        bNode.setColor(Color.RED);
                        RB_DELETE_FIXUP(fixupNode.parent);
                    }

                }
            } else {
                //刪除情景2.2： 替換結點是其父結點的右子結點
                //刪除情景2.2.1： 替換結點的兄弟結點是紅結點
                if (bNode.isRed()) {
                    /*將S設爲黑色 將P設爲紅色 對P進行右旋，獲得情景2.2.2.3 進行情景2.2.2.3的處理*/
                    bNode.setColor(Color.BLACK);
                    fixupNode.parent.setColor(Color.RED);
                    LRotate(fixupNode.parent);
                    RB_DELETE_FIXUP(fixupNode);
                } else { //刪除情景2.2.2： 替換結點的兄弟結點是黑結點
                    //刪除情景2.2.2.1： 替換結點的兄弟結點的左子結點是紅結點，右子結點任意顏色
                    if (bNode.leftChild != null && bNode.leftChild.isRed()) {
                        /*將S的顏色設爲P的顏色 將P設爲黑色 將SL設爲黑色 對P進行右旋*/
                        bNode.color = fixupNode.parent.color;
                        fixupNode.parent.setColor(Color.BLACK);
                        bNode.leftChild.setColor(Color.BLACK);
                        RRotate(fixupNode.parent);
                    } else if (bNode.rightChild != null && bNode.rightChild.isRed()) {//刪除情景2.2.2.2： 替換結點的兄弟結點的左子結點爲黑結點，右子結點爲紅結點
                        /*將S設爲紅色 將SR設爲黑色 對S進行左旋，獲得情景2.2.2.1 進行情景2.2.2.1的處理*/
                        bNode.setColor(Color.RED);
                        bNode.rightChild.setColor(Color.BLACK);
                        LRotate(bNode);
                        RB_DELETE_FIXUP(fixupNode);
                    } else {//刪除情景2.2.2.3：替換結點的兄弟結點的子結點都爲黑結點
                        /*將S設爲紅色 把P做爲新的替換結點 從新進行刪除結點情景處理*/
                        bNode.setColor(Color.RED);
                        RB_DELETE_FIXUP(fixupNode.parent);
                    }
                }
            }
        }
    }

    public static void main(String[] args) {
        //int[] data={8,6,4};
        //int[] data={8,6,9,5,7,3};
        //int[] data={8,6,7};
        //int[] data={8,5,9,4,6,7};
        //int[] data={8,5,9,4,7,6};
        //int[] data = {8, 5, 9, 7, 6};
        //Object[] data = new Object[100];
        //Object[] data = {2, 4, 15, 11, 19, 3, "F", "G", "B", "A", "D", "C", "E", new Person("小王", 22)};
        /*for (int i = 0; i < 100; i++) {
            Random r = new Random();
            int n = r.nextInt(100);
            data[i] = "數據" + n;
            //System.out.println(n);
        }*/
        RBTree rbt = new RBTree();
        int[] data = {2, 4, 15, 11, 19, 3, 12, 14, 16, 9, 13, 17, 7, 8, 5, 1, 18, 6};
        for (int i = 0; i < data.length; i++) {
            if (data[i] == 9) {
                System.out.println();
            }
            System.out.println(data[i]);
            rbt.insertNode(data[i], data[i]);
            rbt.dlrTraversal();
            System.out.println("\n" + rbt.root.data);
        }
        rbt.removen(6);
        /*for (int i = 0; i < data.length; i++) {
            rbt.insertNode(String.valueOf(i), data[i]);
        }*/
        rbt.ldrTraversal();
        System.out.println("\n" + rbt.root.data);
        rbt.dlrTraversal();
        //System.out.println("\n" + rbt.search("0"));
    }
}