自我实现 Java 标准库中TreeMap(简易版)

上一篇文章讲解了 Java 中的 Map 和 Set 的简单使用,这次来分享一个简单版本的 TreeMap.

首先为什么是简单版本,首先 Java 标准库中的 TreeMap 的底层是根据红黑树来实现的,但是由于本人能力有限,红黑树的代码写不出来,所以在自我实现 TreeMap 的过程中使用的是二叉搜索树.

二叉搜索树满足这个特点:每个节点的左孩子的值小于根节点,右孩子的值大于根节点(如果左右孩子都存在的话).

根据二叉搜索树的这个特点就可以在查找的时候很方便.

我主要实现了常用的几个方法

如:get , getOrDefault , put , remove 方法

class TreeNode {
    public int key;
    public int value;
    public TreeNode left;
    public TreeNode right;

    public TreeNode(int key, int value) {
        this.key = key;
        this.value = value;
    }
}

public class BinarySearchTree {
    private TreeNode root = null;

    // 根据 key 的值查找 value
    public Integer get(int key) {
        TreeNode cur = root;
        while (cur != null) {
            if (cur.key > key) {
                cur = cur.left;
            }else if (cur.key < key) {
                cur = cur.right;
            }else {
                return cur.value;
            }
        }
        return null;
    }

    // 根据 key 的值查找 value
    // 如果查找不到，返回默认值
    public Integer getOrDefault(int key, int defaultVal) {
        Integer ret = get(key);
        if (ret == null) {
            return defaultVal;
        }
        return ret;
    }

    // 插入新的键值对
    public void put(int key, int value) {
        // 为空树的时候，直接将这个键值对变为 root 节点
        if (root == null) {
            root = new TreeNode(key, value);
            return;
        }
        TreeNode cur = root;
        TreeNode pre = null;
        while (cur != null) {
            if (key > cur.key) {
                pre = cur;
                cur = cur.right;
            }else if (key < cur.key) {
                pre = cur;
                cur = cur.left;
            }else {
                cur.value = value;
                return;
            }
        }
        if (pre.key > key) {
            pre.left = new TreeNode(key, value);
        }
        if (pre.key < key) {
            pre.right = new TreeNode(key, value);
        }
        return;
    }

    // 根据 key 的值删除对应的键值对
    // 删除成功返回 key 对应的 value ，失败返回 null
    public Integer remove(int key) {
        if (root == null) {
            // 如果 根节点 为 null
            return null;
        }
        TreeNode cur = root;
        TreeNode pre = null;
        // 寻找待删除节点
        while (cur != null) {
            if (key > cur.key) {
                pre = cur;
                cur = cur.right;
            } else if (key < cur.key) {
                pre = cur;
                cur = cur.left;
            } else  {
                // 找到待删除节点
                // 将其删掉
                return _remove1(cur, pre);
            }
        }
        // 循环结束,说明没有找到
        return null;
    }

    // 辅助删除操作,根据上一个节点 pre 删除当前节点 cur
    private Integer _remove(TreeNode cur, TreeNode pre) {
        Integer ret = cur.value;
        if (cur == root) {
            // 待删除节点是根节点
            if (cur.left == null) {
                // 根节点的左子树为空
                root = root.right;
                return ret;
            }else if (cur.right == null) {
                // 根节点的右子树为空
                root = root.left;
                return ret;
            }else {
                // 根节点的左右子树都不为空
                // 找到根节点左子树的最右边节点,覆盖根节点,再将这个节点删掉
                TreeNode goat = root.left;
                if (goat.right == null) {
                    // 如果左子树没有右孩子
                    root.key = goat.key;
                    root.value = goat.value;
                    root.left = goat.left;
                    return ret;
                }else {
                    // 如果左子树有右孩子
                    pre = null;
                    while (goat.right != null) {
                        pre = goat;
                        goat = goat.right;
                    }
                    // 找到了,下面覆盖根节点
                    root.key = goat.key;
                    root.value = goat.value;
                    pre.right = goat.left;
                    return ret;
                }
            }
        }else {
            // 待删除节点不是根节点
            if (pre.left == cur) {
                // 待删除节点是前一个节点的左孩子
                if (cur.left == null) {
                    // 待删除节点没有左孩子
                    pre.left = cur.right;
                    return ret;
                }else if (cur.right == null) {
                    // 待删除节点没有右孩子
                    pre.left = cur.left;
                    return ret;
                }else {
                    // 待删除节点既有左孩子也有右孩子
                    TreeNode goat = cur.left;
                    if (goat.right == null) {
                        cur.key = goat.key;
                        cur.value = goat.value;
                        cur.left = goat.left;
                        return ret;
                    }else {
                        while (goat.right != null) {
                            pre = goat;
                            goat = goat.right;
                        }
                        cur.key = goat.key;
                        cur.value = goat.value;
                        pre.right = goat.left;
                        return ret;
                    }
                }
            }else if (pre.right == cur) {
                // 待删除节点是前一个节点的右孩子
                if (cur.left == null) {
                    // 待删除节点没有左孩子
                    pre.right = cur.right;
                    return ret;
                }else if (cur.right == null) {
                    // 待删除节点没有右孩子
                    pre.right = cur.left;
                    return ret;
                }else {
                    // 待删除节点既有左节点又有右节点
                    TreeNode goat = cur.left;
                    if (goat.right == null) {
                        cur.key = goat.key;
                        cur.value = goat.value;
                        cur.left = goat.left;
                        return ret;
                    }else {
                        while (goat.right != null) {
                            pre = goat;
                            goat = goat.right;
                        }
                        cur.key = goat.key;
                        cur.value = goat.value;
                        pre.right = goat.left;
                        return ret;
                    }
                }
            }
        }
        return ret;
    }

    // 辅助删除操作,根据上一个节点 pre 删除当前节点 cur
    private Integer _remove1(TreeNode cur, TreeNode pre) {
        Integer ret = cur.value;
       if (cur.left == null) {
           // 待删除节点的左子树为空
           if (root == cur) {
               // 如果待删除节点是根节点
               root = root.right;
           }else {
               // 不是根节点
               if (pre.left == cur) {
                   // 待删除节点是前一个节点的左孩子
                   pre.left = cur.right;
               }else {
                   // 待删除节点是前一个节点的右孩子
                   pre.right = cur.right;
               }
           }
           return ret;
       }else if (cur.right == null) {
           // 待删除节点的右子树为空
           if (cur == root) {
               // 待删除节点是根节点
               root = root.left;
           }else {
               // 不是根节点
               if (pre.left == cur) {
                   // 待删除节点是前一个节点的左孩子
                   pre.left = cur.left;
               }else {
                   // 待删除节点是前一个节点的右孩子
                   pre.right = cur.left;
               }
           }
           return ret;
       }else {
           // 待删除节点的左右孩子都不为空
           TreeNode goat = cur.left;
           if (goat.right == null) {
               cur.key = goat.key;
               cur.value = goat.value;
               cur.left = goat.left;
           }else {
               while (goat.right != null) {
                   pre = goat;
                   goat = goat.right;
               }
               cur.key = goat.key;
               cur.value = goat.value;
               pre.right = goat.left;
           }
           return ret;
       }
    }
}

实现这个类最复杂的工作是删除操作 remove ,因为根据节点的情况不同删除要区分很多种情况.

根据上一个节点 pre 删除当前节点 cur 我写了两个方法,分别是_remove 和 _remove1 

其中 _remove 方法由于分类不当导致冗余度很大,

_remove 中,我是按照这样的逻辑来分类的:(反面教材不要学习)

1 待删除节点是根节点

1.1 根节点的左子树为空

1.2 根节点的右子树为空

1.3 根节点的左右子树都不为空

1.3.1 根节点的左子树没有右孩子

1.3.2 根节点的左子树有右孩子

2 待删除节点不是根节点

2.1 待删除节点是前一个节点的左孩子

2.1.1 待删除节点没有左孩子

2.1.2 待删除节点没有右孩子

2.1.3 待删除节点既有左孩子也有右孩子

2.1.3.1 待删除节点的左孩子没有右孩子

2.1.3.2 待删除节点的左孩子有右孩子

2.2 待删除节点是前一个节点的右孩子

2.2.1 待删除节点没有左孩子

2.2.2 待删除节点没有右孩子

2.2.3 待删除节点既有左孩子也有右孩子

2.2.3.1 待删除节点的左孩子没有右孩子

2.2.3.2 待删除节点的左孩子有右孩子

这种分类方式比较符合人类逻辑,但是导致了代码冗余量巨大,其中很多操作的代码一模一样.

实际上 当我们进行 if else 分类的时候如果处理方式一样,那么就可以不用分类,例如当待删除节点的左右孩子都存在时,不管它是不是根节点我们要执行的操作都是相同的.

所以代码的逻辑要从现实意义出发,向上面的分类看似很缜密其实落实下来很不容易.

下面我们来看 _remove1 方法的分类逻辑

1 待删除的节点的左子树为空

1.1 待删除节点是根节点

1.2 待删除节点不是根节点

1.2.1 待删除节点是前一个节点的左孩子

1.2.2 待删除节点是前一个节点的右孩子

2 待删除的节点的右子树为空

2.1 待删除节点是根节点

2.2 待删除节点不是根节点

2.2.1 待删除节点是前一个节点的左孩子

2.2.2 待删除节点是前一个节点的右孩子

3 待删除的节点的左右子树都不为空

3.1 待删除节点的左孩子没有右孩子

3.2 待删除节点的左孩子有右孩子

这样的分类就简单很多

为什么待删除节点的左右孩子都存在的时候不考虑是不是根节点的情况?

因为我们这里的处理方式是找到待删除节点的左子树的最右侧节点,将它的值复制到待删除节点位置,再将这个节点删除,和它是不是根节点无关.这里的处理逻辑一样我们就没有必要再分了.

下面给出测试代码:

    public static void main(String[] args) {
        BinarySearchTree bt = new BinarySearchTree();
        bt.put(1,1);
        bt.put(2,2);
        bt.put(3,3);
        bt.put(4,4);
        bt.put(5,5);
        bt.put(6,6);
        System.out.println(bt.get(4));
        System.out.println(bt.getOrDefault(8,150));
        bt.remove(3);
        bt.remove(4);
        bt.remove(5);
        System.out.println(bt.get(3));
        System.out.println(bt.get(4));
        System.out.println(bt.get(5));
        System.out.println(bt.get(6));

    }

测试结果如下: