1.4.3.2 RB Tree Insertion Implementation

完整的原始碼點我

這邊是insert的部分:

public void insert(int id, int data) {
        // Create a new node
        RbNode newNode = new RbNode(id, data, true);

        if (root == null) {
            root = newNode;
        } else {
            // Find the insert location
            RbNode current = root;
            RbNode parent = null;

            while (true) {
                parent = current;
                if (id < current.getId()) {
                    current = current.getLeft();
                    // If no left
                    if (current == null) {
                        // Modify related node's attribute
                        parent.setLeft(newNode);
                        newNode.setParent(parent);
                        break;
                    }
                } else {
                    current = current.getRight();
                    // If no right
                    if (current == null) {
                        // Modify related node's attribute
                        parent.setRight(newNode);
                        newNode.setParent(parent);
                        break;
                    }
                }
            }
        }

        balanceAfterInsert(newNode);
    }

這邊是平衡的部分:

private void balanceAfterInsert(RbNode currentNode) {
        /*
         * Case 1: If the node is root, this will violate the rule
         * that the root should be black, so change the node to black.
         */
        if (currentNode.getParent() == null) {
            currentNode.setRed(false);
            root = currentNode;
        } else
        /*
         * Case 2: If node P is black, do nothing since it already match the rule.
         */
        if (!currentNode.getParent().isRed()) {
            // Do nothing.
        } else if (currentNode.getParent().isRed()) {
            RbNode gNode = currentNode.getParent().getParent();
            RbNode uNode = null;
            if (gNode != null) {
                uNode = (gNode.getLeft() == currentNode.getParent()) ? gNode.getRight() : gNode.getLeft();
            }
            /*
             * Case 3: If node P and node U are all red, modify node G to red and modify both node P
             * and node U to black, then set the node G as current node and perform
             * balanceAfterInsert
             * operation on it.
             */
            if (uNode != null && uNode.isRed()) {
                gNode.setRed(true);
                uNode.setRed(false);
                currentNode.getParent().setRed(false);
                currentNode = gNode;
                balanceAfterInsert(currentNode);
            } else if (uNodeIsBlackOrNull(uNode)) {
                /*
                 * Case 4: If node P is red and node U is black, and the inserted node is the left
                 * of node P, and the node P is the left of node G, then modify the node P to black,
                 * and modify the node G to red. Finally, perform right rotate on node G and
                 * balanceAfterInsert
                 * on current node.
                 */
                if (currentNode == currentNode.getParent().getLeft()
                        && (gNode != null && currentNode.getParent() == gNode.getLeft())) {
                    currentNode.getParent().setRed(false);
                    gNode.setRed(true);
                    rightRotate(gNode);
                    balanceAfterInsert(currentNode);
                } else
                /*
                 * Case 5: If node P is red and node U is black, and the inserted node is the right
                 * of node P, and the node P is the right of node G, then modify the node P to
                 * black, and modify the node G to red. Finally, perform left rotate on node G and
                 * balanceAfterInsert on current node.
                 */
                if (currentNode == currentNode.getParent().getRight()
                        && (gNode != null && currentNode.getParent() == gNode.getRight())) {
                    currentNode.getParent().setRed(false);
                    gNode.setRed(true);
                    leftRotate(gNode);
                    balanceAfterInsert(currentNode);
                } else
                /*
                 * Case 6: If node P is red and node U is black, and the inserted node is the right
                 * of node P, and the node P is the left of node G, then let the node P be the new
                 * current node. Finally, perform left rotate on the current node and
                 * balanceAfterInsert on it
                 */
                if (currentNode == currentNode.getParent().getRight()
                        && (gNode != null && currentNode.getParent() == gNode.getLeft())) {
                    RbNode oldParent = currentNode.getParent();
                    leftRotate(oldParent);
                    balanceAfterInsert(oldParent);
                } else
                /*
                 * Case 7:If node P is red and node U is black, and the inserted node is the left of
                 * node P, and the node P is the right of node G, then let the node P be the new
                 * current node. Finally perform right rotate on the current node and
                 * balanceAfterInsert on it
                 */
                if (currentNode == currentNode.getParent().getLeft()
                        && (gNode != null && currentNode.getParent() == gNode.getRight())) {
                    RbNode oldParent = currentNode.getParent();
                    rightRotate(oldParent);
                    balanceAfterInsert(oldParent);
                }
            }
        }
    }

    private boolean uNodeIsBlackOrNull(RbNode uNode) {
        return (uNode == null || !uNode.isRed());
    }

這邊是左旋轉的部分:

private void leftRotate(RbNode pivot) {
        RbNode oldRight = pivot.getRight();
        RbNode leftOfOldRight = null;
        if (oldRight != null) {
            leftOfOldRight = oldRight.getLeft();
        }

        if (pivot.getParent() != null) {
            // Determine it's the left ot right
            boolean isLeft = (pivot.getParent().getLeft() == pivot);
            if (isLeft) {
                pivot.getParent().setLeft(oldRight);
            } else {
                pivot.getParent().setRight(oldRight);
            }

            if (oldRight != null) {
                oldRight.setParent(pivot.getParent());
            }
        } else {
            oldRight.setParent(null);
            oldRight.setRed(false);
            root = oldRight;
        }

        if (oldRight != null) {
            oldRight.setLeft(pivot);
        }
        pivot.setParent(oldRight);

        pivot.setRight(leftOfOldRight);
        if (leftOfOldRight != null) {
            leftOfOldRight.setParent(pivot);
        }
    }

這邊是右旋轉的部分:

private void rightRotate(RbNode pivot) {
        RbNode oldLeft = pivot.getLeft();
        RbNode rightOfOldLeft = null;
        if (oldLeft != null) {
            rightOfOldLeft = oldLeft.getRight();
        }
        if (pivot.getParent() != null) {
            // Determine it's the left ot right
            boolean isLeft = (pivot.getParent().getLeft() == pivot);
            if (isLeft) {
                pivot.getParent().setLeft(oldLeft);
            } else {
                pivot.getParent().setRight(oldLeft);
            }

            if (oldLeft != null) {
                oldLeft.setParent(pivot.getParent());
            }
        } else {
            oldLeft.setParent(null);
            oldLeft.setRed(false);
            root = oldLeft;
        }

        if (oldLeft != null) {
            oldLeft.setRight(pivot);
        }
        pivot.setParent(oldLeft);

        pivot.setLeft(rightOfOldLeft);
        if (rightOfOldLeft != null) {
            rightOfOldLeft.setParent(pivot);
        }
    }