带你手写红黑树——500行代码

手写一个红黑树

既然红黑树如此的不友好，那么我就来手写一个红黑树，来细致地探究一下该如何实现一个属于自己的红黑树,这样自己也好有个底^_。

500+行代码实现一个红黑树：

树的设计

public class RBTree<Key extends Comparable<Key>, Value> {
    /****************颜色常量***************************/
    private static final boolean RED = true;
    private static final boolean BLACK = false;
    /****************旋转类型常量*********************/
    private static final int LEFT_RIGHT 	= -2;
    private static final int LEFT_LEFT 		= -1;
    private static final int NONE 			=  0;
    private static final int RIGHT_LEFT 	=  1;
    private static final int RIGHT_RIGHT	=  2;
    private TreeNode root;// 根节点
}

红黑树内部有一个可以比较大小的的 Key字段，用于作为红黑树节点 TreeNode(RBTree中的一个内部类)的键值，存在一个 Value的数据项，和一些需要使用到的常量。

TreeNode树节点的数据结构

private class TreeNode {
    Key key;		// 键值
    Value value;	// 数据
    TreeNode left;	// 左节点
    TreeNode right;	// 右节点
    TreeNode parent;// 父节点
    int amount;		// 子树节点数量
    boolean color;	// 节点颜色
}

左旋算法实现

动图：

代码：

/***
 * @param node 旋转中心节点
 * @return TreeNode 旋转后的中心节点
 */
private TreeNode rotateLeft(TreeNode node) {
    // 第一步：判断是否能左旋,即是判断旋转中心节点的右节点是否存在
    if (node.right == null) {
        return node;
    }
    TreeNode par = node.parent;
    TreeNode right = node.right;
    // 第二步：将node的父节点的引用赋值给node.right的<code>parent</code>属性
    right.parent = par;
    if (par != null) {
        if (node == par.left) {
            par.left = right;
        } else {
            par.right = right;
        }
    }
    // 第三步：将node的右节点设置为right的左节点
    node.right = right.left;
    node.amount = getAmount(node.right) + getAmount(node.left) + 1;
    // 第4步：node右节点的父节点设置为node
    if (node.right != null) {
        node.right.parent = node;
    }
    // 第5步：设置right的左节点为node
    right.left = node;
    right.amount = getAmount(right.left) + getAmount(right.right) + 1;
    node.parent = right;
    return right;
}

右旋算法实现

动图

代码

/***
 * @param node 旋转中心节点
 * @return TreeNode 旋转后的中心节点
 */
private TreeNode rotateRight(TreeNode node) {
        // 第一步：判断是否能右旋,即是判断旋转中心节点的右节点是否存在
//        HashMap
        if (node.left == null) {
            return node;
        }
        TreeNode par = node.parent;
        TreeNode left = node.left;
        // 第二步：将node的父节点的引用赋值给node.left<code>parent</code>属性
        left.parent = par;
        if (par != null) {
            if (node == par.left) {
                par.left = left;
            } else {
                par.right = left;
            }
        }
        // 第三步：将node的左节点设置为left的右节点
        node.left = left.right;
        node.amount = getAmount(node.right) + getAmount(node.left) + 1;
        // 第4步：node左节点的父节点设置为node
        if (node.left != null) {
            node.left.parent = node;
        }
        // 第5步：设置left的右节点为node
        left.right = node;
        left.amount = getAmount(left.left) + getAmount(left.right) + 1;
        node.parent = left;
        return left;
    }

put()方法实现

/**
 * put 添加(key冲突，默认执行替换)
 */
public TreeNode put(Key key, Value value) {
    TreeNode target = new TreeNode(key, value, null, null, null, 1, RED);
    return this.putNode(root, target, true);
}

/**
 * 向树中添加一个由 key - value 包装的节点，并自动进行红黑树平衡
 * put 添加(key冲突，由参数 conflict_action 决定是否替换，true 替换 ， false 不替换)
 */
public TreeNode put(Key key, Value value, boolean conflict_action) {
    TreeNode target = new TreeNode(key, value, null, null, null, 1, RED);
    return this.putNode(root, target, conflict_action);
}

其实这是在套用 putNode()方法：

/**
 * 向树中某个节点处添加一个由 key - value 包装的节点，并自动进行红黑树平衡
 * @param node 当前节点
 * @param target 待插入节点
 * @param conflict_action key冲突处理方案
 * @return 待插入节点
 */
private TreeNode putNode(TreeNode node, TreeNode target, boolean conflict_action) {
    if (node == null) {//根节点没有初始化
        target.color = BLACK;
        return root = target;
    }
    int flag = target.key.compareTo(node.key);//判断插入子树位置
    if (flag == 0) {//key值冲突，解决冲突
        if (conflict_action) {
            node.value = target.value;
        }
        return target;
    } else if (flag < 0) {//应该插入到左子树
        if (node.left == null) {//当前节点的左子树为 null ,直接插入
            node.left = target;
            target.parent = node;
            changeAmountToRoot(node, 1);
            rotate(target);//旋转平衡
            return target;
        }
        return putNode(node.left, target, conflict_action);
    } else {
        if (node.right == null) {
            node.right = target;
            target.parent = node;
            changeAmountToRoot(node, 1);
            rotate(target);
            return target;
        }
        return putNode(node.right, target, conflict_action);
    }
}

在这个方法里面最关键的还是 rotate()方法：

自平衡方法——rotate()方法

/**
 * 根据给定的节点，进行需要的旋转，包括颜色的更改
 * 改变树结构符合红黑树标准
 */
private void rotate(TreeNode node) {
    int kind = kindOfRotate(node);//得到旋转的类型
    TreeNode p = node.parent;
    if (p == null) {//说明待旋转的节点为根节点，直接返回即可
        node.color = BLACK;
        root = node;
        return;
    }
    TreeNode gp = p.parent;
    if (gp == null) {//第二层子树节点，不必旋转
        return;
    }
    switch (kind) {
        case LEFT_LEFT: {// LL
            gp.color = RED;
            p.color = BLACK;
            // 如果叔叔节点存在且为红色节点
            if(isRed(gp.right)){
                gp.right.color = BLACK;// 将颜色改为黑色
                rotate(gp);// 以gp为当前节点继续操作，一直递归到根节点
            }else {
                // 叔叔节点为黑色节点或者不存在
                // 以gp右旋后的的顶点为目标节点
                // 继续递归直到根节点
                rotate(rotateRight(gp));
            }
            break;
        }
        case LEFT_RIGHT: {// LR
            rotateLeft(p);// 以父亲节点左旋，使其成为 LL
            gp.color = RED;
            node.color = BLACK;
            if(isRed(gp.right)){
                gp.right.color = BLACK;
                rotate(gp);
            }else {
                rotate(rotateRight(gp));
            }
            break;
        }
        case RIGHT_LEFT: {
            rotateRight(p);
            gp.color = RED;
            node.color = BLACK;
            if(isRed(gp.left)){
                gp.left.color = BLACK;
                rotate(gp);
            }else {
                rotate(rotateLeft(gp));
            }
            break;
        }
        case RIGHT_RIGHT: {
            gp.color = RED;
            p.color = BLACK;
            if(isRed(gp.left)){
                gp.left.color = BLACK;
                rotate(gp);
            }else {
                rotate(rotateLeft(gp));
            }
            break;
        }
    }
}

源码获取方式

至此，只是展示了一部分代码，如果需要完整的源代码，请移步个人公众号【然Coder】，输入【红黑树源码】即可获得源代码文件。

运行效果