本文记录伯克利的数据结构与算法课程 -- CS61B 21spring中Lab7的实现思路。

该Lab是实现一个基于二叉搜索树(BST)的Map。
完整代码可以参考我的GitHub仓库。

BSTMap需要实现的接口

/** Removes all of the mappings from this map. */
void clear();
/* Returns true if this map contains a mapping for the specified key. */
boolean containsKey(K key);
/* Returns the value to which the specified key is mapped, or null if this map contains no mapping for the key. */
V get(K key);
/* Returns the number of key-value mappings in this map. */
int size();
/* Associates the specified value with the specified key in this map. */
void put(K key, V value);
/* Returns a Set view of the keys contained in this map. Not required for Lab 7.  */
Set<K> keySet();
/* Removes the mapping for the specified key from this map if present.Challenging! Return null if nothing. 
*/
V remove(K key);
/* Removes the entry for the specified key only if it is currently mapped to the specified value.*/
V remove(K key, V value);
/* Return an iterator for the key.  **/
Iterator<T> iterator();

BSTMap实现

二叉搜索树BST作为一种经典的递归定义数据结构，其实现非常具有典型意义，是锻炼递归的好机会。

定义

BSTMap需要定义泛型，且Key的泛型需要支持比较，即继承Comparable接口，除此之外，按照文档要求，BSTMap的接口定义在接口 Map61B<K,V>中，此接口继承了 Iterable<K>。

BST拥有一个内部类BSTNode，我参考普林斯顿大学算法书对BST的定义，在BSTNode除了设置常规的key、value和left、right指针外，还维护了一个size。而BSTNode只需要维护一个根节点root即可，代码如下。

public class BSTMap<K extends Comparable<K>,V> implements Map61B<K ,V>{

    private BSTNode root;

    private class BSTNode{
        private K key;
        private V val;
        private BSTNode left,right;
        private int size;

        public BSTNode(K key, V val, int size){
            this.key = key;
            this.val = val;
            this.size = size;
        }
    }

    public BSTMap() {

    }
}

接口实现

由于BSTMap是一个递归定义的数据结构，几乎所有的接口我都定义了一个私有的辅助函数。例如，containsKey(K key)会有一个辅助函数，containsKey(BSTNode node , K key)。

所以接口的实现套路基本相同，先判断是否为空，即递归出口，然后调用CompareTo比较，根据比较结果，递归调用左右子树。

size的实现

size只负责返回节点的size属性，不负责改变。

这里设计的妙处就在于，size是由节点本身维护的属性，节点失去引用了，size自然就减小了。我们只需要从叶子节点往上更新一遍size即可，不需要真正对size进行操作！

    @Override
    public int size() {
        return size(root);
    }

    private int size(BSTNode node){
        if(node == null){
            return 0;
        }
//        将size的改变放到add中，维护size，提高速度
//        return size(node.left) + size(node.right) + 1;
        return node.size;
    }

put的实现

put需要判断value是否为空，如果为空就执行remove操作即可，这是一个小细节。

最后记得更新一下size，更新size的逻辑和remove相同，左子树+右子树+1即可。

    @Override
    public void put(K key, V value) {
        if(key == null) {
            throw new IllegalArgumentException("key should not be null");
        }
//        如果value为空，执行删除操作
        if(containsKey(key) && value == null){
            remove(key);
        }
        root = put(root,key,value);
    }

    private BSTNode put(BSTNode node, K key, V value){
        if(node == null) {
            return new BSTNode(key,value,1);
        }
        int cmp = key.compareTo(node.key);
        if(cmp == 0) {
            node.val = value;
        }else if(cmp < 0){
            // key < node.key , put it left
            node.left =  put(node.left,key,value);
        }else{
            node.right = put(node.right,key,value);
        }
        node.size = 1 + size(node.left) + size(node.right);
        return node;
    }

keySet实现

keySet返回所有key的一个Set集合。这里我们采用HashSet。因为一样需要遍历所有节点，所以这里也设置了一个辅助函数，做一个先序遍历就可以。

    @Override
    public Set<K> keySet() {
        if(root == null){
            return null;
        }
        HashSet<K> ks = new HashSet<>();
        addKey(root,ks);
        return ks;
    }
//    辅助函数，递归添加set
    private void addKey(BSTNode node, Set<K>set){
        if(node == null){
            return ;
        }
        // 先序遍历
        set.add(node.key);
        addKey(node.left,set);
        addKey(node.right,set);
    }

remove实现

删除逻辑是比较复杂的部分内容，我们需要分类讨论。

  1. 如果只有左子树或者右子树：直接将儿子赋给自己。
  2. 如果同时有左和右，找到右子树的最小元素赋给自己，删除右子树的最小元素。

可见：要实现删除，需要构造【找到最小】和【删除最小】两个辅助函数。

首先定义一个删除最小的函数：实现思路是递归查找左子树，找到最小的节点，当某个节点的左儿子是叶子节点的时候，将让左儿子指向右孙子。

需要注意的是，需要更新一次节点。

    /** Remove min node in BSTMap */
    private BSTNode removeMin(BSTNode node){
//        找到最小的节点，当最小的时候，将让左儿子指向右孙子
        if(node.left == null){
            return node.right;
        }
        node.left =  removeMin(node.left);
//        注意减小size，实际上，改变指针的指向后，size函数就不会统计对应的大小了。
//        这里主要是更新父节点的size
        node.size = size(node.left) + size(node.right) + 1;
        return node;
    }

再实现一个找到最小节点的函数：也很简单，一直找左子树即可。

    /** Return the smallest node */
    private BSTNode min(BSTNode node){
        if(node.left == null) {
            return node;
        }
        return min(node.left);
    }

最后就是删除函数和他的辅助函数了，按照上面的分析思路实现即可。

 @Override
    public V remove(K key) {
        if(key == null) {
            throw new IllegalArgumentException("key should not be null.");
        }
        if(!containsKey(key)){
            return null;
        }
        V val = get(key);
//      将查找的逻辑放到帮助函数内
        root = remove(root,key);
        return val;
    }

    private BSTNode remove(BSTNode node, K key){
        if(node == null){
            return null;
        }
        int cmp = key.compareTo(node.key);
        if(cmp < 0){
            // key < node.key , find in left
            node.left = remove(node.left,key);
        }else if(cmp > 0){
            node.right = remove(node.right,key);
        }else{
            // 分类讨论
            if(node.right == null) return node.left;
            if(node.left == null) return node.right;
            // 临时变量指向node，然后node指向右子树的最小节点
            BSTNode tmp = node;
            node = min(node.right);
            node.right = removeMin(tmp.right);
            node.left = tmp.left;
        }
//        更新size
        node.size = size(node.left) + size(node.right)  + 1;
        return node;
    }

iterator实现

最后是iterator的实现，实际上只需要返回 keySet.iterator()即可。

参考资料

{% referfrom '[1]',' UCBerkeley CS61B Lab7 Document','https://sp21.datastructur.es/materials/lab/lab7/lab7' %}

{% referfrom '[2]',' Princeton Algorithms 4th edtion','https://algs4.cs.princeton.edu/32bst/BST.java.html' %}