如何在 Java 中實現二叉搜尋樹

之一Yo發表於2022-03-24

二叉搜尋樹

二叉搜尋樹結合了無序連結串列插入便捷和有序陣列二分查詢快速的特點,較為高效地實現了有序符號表。下圖顯示了二叉搜尋樹的結構特點(圖片來自《演算法第四版》):

二叉搜尋樹的結構

可以看到每個父節點下都可以連著兩個子節點,鍵寫在節點上,其中左邊的子節點的鍵小於父節點的鍵,右節點的鍵大於父節點的鍵。每個父節點及其後代節點組成了一顆子樹,根節點及其後代節點則組成了完整的二叉搜尋樹。

在程式碼層面看來,就是每個節點物件中包含另外兩個子節點的指標,同時包含一些要用到的資料,比如鍵值對和方便後續操作的整課子樹的節點數量。

private class Node {
    int N = 1;
    K key;
    V value;
    Node left;
    Node right;

    public Node(K key, V value) {
        this.key = key;
        this.value = value;
    }
}

上述程式碼實現了一個節點類,這個類是二叉搜尋樹類 BinarySearchTree 的內部類,使用者無需知道這個節點類的存在,所以訪問許可權宣告為 private。

有序符號表的 API

先來看下無序符號表的 API,這些方法宣告瞭無序符號表的基本操作,包括插入、查詢和刪除,為了方便符號表的迭代,介面中還有 Iterable<K> keys() 方法用於 foreach 迴圈:

package com.zhiyiyo.collection.symboltable;

public interface SymbolTable<K, V>{
    void put(K key, V value);
    V get(K key);
    void delete(K key);
    boolean contains(K key);
    boolean isEmpty();
    int size();
    Iterable<K> keys();
}

接下來是有序符號表的 API,其中每個節點的鍵必須實現了 Comparable 介面:

package com.zhiyiyo.collection.symboltable;

public interface OrderedSymbolTable<K extends Comparable<K>, V> extends SymbolTable<K, V>{
    /**
     * 獲取符號表中最小的鍵
     * @return 最小的鍵
     */
    K min();

    /**
     * 獲取符號表中最大的鍵
     * @return 最大的鍵
     */
    K max();

    /**
     * 獲取小於或等於 key 的最大鍵
     * @param key 鍵
     * @return 小於或等於 key 的最大鍵
     */
    K floor(K key);

    /**
     * 獲取大於或等於 key 的最小鍵
     * @param key 鍵
     * @return 大於或等於 key 的最小鍵
     */
    K ceiling(K key);

    /**
     * 獲取小於或等於 key 的鍵數量
     * @param key 鍵
     * @return 小於或等於 key 的鍵數量
     */
    int rank(K key);

    /**
     * 獲取排名為 k 的鍵,k 的取值範圍為 [0, N-1]
     * @param k 排名
     * @return 排名為 k 的鍵
     */
    K select(int k);

    /**
     * 刪除最小的鍵
     */
    void deleteMin();

    /**
     * 刪除最大的鍵
     */
    void deleteMax();

    /**
     * [low, high] 區間內的鍵數量
     * @param low 最小的鍵
     * @param high 最大的鍵
     * @return 鍵數量
     */
    int size(K low, K high);

    /**
     * [low, high] 區間內的所有鍵,升序排列
     * @param low 最小的鍵
     * @param high 最大的鍵
     * @return 區間內的鍵
     */
    Iterable<K> keys(K low, K high);
}

實現二叉搜尋樹

二叉搜尋樹類

類的基本結構如下述程式碼所示,可以看到只需用一個根節點 root 即可代表一整棵二叉搜尋樹:

public class BinarySearchTree<K extends Comparable<K>, V> implements OrderedSymbolTable<K, V> {
    private Node root;

    private class Node{
        ...
    }

    ...
}

查詢

從根節點出發,拿著給定的鍵 key 和根節點的鍵進行比較,會出現以下三種情況:

  • 根節點的鍵大於 key,接著去根節點的左子樹去查詢;
  • 根節點的鍵小於 key,接著去根節點的右子樹去查詢;
  • 根節點的鍵等於 key,返回根節點的值

當我們去左子樹或者右子樹查詢時,只需將子樹的根節點視為新的根節點,然後重複上述步驟即可。如果找到最後都沒找到擁有和 key 相等的鍵的節點,返回 null 即可。在《演算法第四版》中使用遞迴實現了上述步驟,這裡換為迭代法:

@Override
public V get(K key) {
    Node node = root;
    while (node != null) {
        int cmp = node.key.compareTo(key);
        // 到左子樹搜尋
        if (cmp > 0) {
            node = node.left;
        }
        // 到右子樹搜尋
        else if (cmp < 0) {
            node = node.right;
        } else {
            return node.value;
        }
    }
    return null;
}

插入

將鍵值對放入二叉搜尋樹時會發生兩種情況:

  • 二叉搜尋樹中已經包含了擁有該鍵的節點,這時需要更新節點的值
  • 二叉搜尋樹中沒有包含擁有該鍵的節點,這時需要建立一個新的節點

所以在插入的時候要從根節點出發,比較根節點的鍵和給定的 key 之間的大小關係,和查詢相似,比較會有三種情況發生:

  • 根節點的鍵大於 key,接著去根節點的左子樹去查詢;
  • 根節點的鍵小於 key,接著去根節點的右子樹去查詢;
  • 根節點的鍵等於 key,直接更新根節點的值

如果找到最後都沒能找到那個擁有相同 key 的節點,就需要建立一個新的節點,把這個節點,接到子樹的根節點上,用迭代法實現上述過程的程式碼如下所示:

@Override
public put(K key, V value){
    if (root == null) {
        root = new Node(key, value);
        return;
    }

    Node node = root;
    Node parent = root;
    int cmp = 0;

    while (node != null){
        parent = node;
        cmp = node.key.compareTo(key);
        // 到左子樹搜尋
        if (cmp > 0){
            node = node.left;
        }
        // 到右子樹搜尋
        else if (cmp < 0){
            node = node.right;
        } else {
            node.value = value;
            return;
        }
    }

    // 新建節點並接到父節點上
    if (cmp > 0) {
        parent.left = new Node(key, value);
    } else{
        parent.right = new Node(key, value);
    }
}

可以看到上述過程用了兩個指標,一個指標 node 用於探路,一個指標 parent 用於記錄子樹的根節點,不然當 node 為空時我們是找不到他的父節點的,也就沒法把新的節點接到父節點上。

上述程式碼有個小問題,就是我們新建節點之後沒辦法更新這一路上所經過的父節點的 N,也就是每一顆子樹的節點數。怎麼辦呢,要麼用一個容器儲存一下經過的父節點,要麼老老實實用遞迴,這裡選擇用遞迴。遞迴的想法很直接:

  • 如果根節點的鍵大於 key,就把鍵值對插到根節點的左子樹;
  • 如果根節點的鍵小於 key,就把鍵值對插到根節點的右子樹;
  • 如果根節點的鍵等於 key,直接更新根節點的值

別忘了,使用遞迴的原因是我們要更新父節點的 N,所以遞迴的返回值應該是更新後的子樹根節點,所以就有了下述程式碼:

@Override
public void put(K key, V value) {
    root = put(root, key, value);
}

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

private int size(Node node) {
    return node == null ? 0 : node.N;
}

最小/大的鍵

從根節點出發,一路向左,鍵會是一個遞減的序列,當我們走到整棵樹的最左邊,也就是 leftnull 的那個節點時,我們就已經找到了鍵最小的節點。上述過程的迭代法程式碼如下:

@Override
public K min() {
    if (root == null) {
        return null;
    }

    Node node = root;
    while (node.left != null) {
        node = node.left;
    }

    return node.key;
}

查詢最大鍵的節點過程和上述過程類似,只是我們這次得向右走,直到找到 rightnull 的那個節點:

@Override
public K max() {
    if (root == null) {
        return null;
    }

    Node node = root;
    while (node.right != null) {
        node = node.right;
    }

    return node.key;
}

演算法書中給出的 min() 實現程式碼是用遞迴實現的,因為在刪除節點時會用到。遞迴的過程就是一直朝左子樹走的的過程,直到遇到一個節點沒有左子樹為止,然後返回該節點即可。

@Override
public K min() {
    if (root == null) {
        return null;
    }

    return min(root).key;
}

private Node min(Node node) {
    if (node.left == null) return node;
    return min(node.left);
}

小於等於 key 的最大鍵/大於等於 key 的最小鍵

從根節點出發,拿著根節點的的鍵和 key 進行比較,會出現三種情況:

  • 如果根節點的鍵大於 key,說明擁有小於或等於 key 的鍵的節點可能在左子樹上(也可能找不到);
  • 如果根節點的鍵小於 key,這時候先記住根節點,由於根節點的右子樹上可能存在鍵更接近但不大於 key 的節點,所以還得去右子樹看看,如果右子樹沒沒找到滿足條件的節點,這時候的根節點的鍵就是小於等於 key 的最大鍵了;
  • 如果根節點的鍵等於 key,直接返回根節點的鍵
@Override
public K floor(K key) {
    if (root == null) {
        return null;
    }

    Node node = root;
    Node candidate = root;
    while (node != null) {
        int cmp = node.key.compareTo(key);
        if (cmp > 0) {
            node = node.left;
        } else if (cmp < 0) {
            candidate = node;
            node = node.right;
        } else {
            return node.key;
        }
    }

    return candidate.key.compareTo(key) <= 0 ? candidate.key : null;
}

《演算法第四版》中給出了一個示例圖,可以更直觀地看到上述查詢過程:

floor

查詢大於等於 key 的最小鍵的方法和上述過程很像,拿著根節點的的鍵和 key 進行比較,會出現三種情況:

  • 如果根節點的鍵小於 key,說明擁有大於或等於 key 的鍵的節點可能在右子樹上(也可能找不到);
  • 如果根節點的鍵大於 key,這時候先記住根節點,由於根節點的左子樹上可能存在鍵更接近但不小於 key 的節點,所以還得去左子樹看看,如果左子樹沒沒找到滿足條件的節點,這時候的根節點的鍵就是大於等於 key 的最小鍵了;
  • 如果根節點的鍵等於 key,直接返回根節點的鍵
@Override
public K ceiling(K key) {
    if (root == null) {
        return null;
    }

    Node node = root;
    Node candidate = root;
    while (node != null) {
        int cmp = node.key.compareTo(key);
        if (cmp < 0) {
            node = node.right;
        } else if (cmp > 0) {
            candidate = node;
            node = node.left;
        } else {
            return node.key;
        }
    }

    return candidate.key.compareTo(key) >= 0 ? candidate.key : null;
}

根據排名獲得鍵

假設一棵二叉搜尋樹中有 N 個節點,那麼節點的鍵排名區間就是 [0, N-1],也就是說,key 的排名可以看做小於 key 的鍵的個數。所以我們應該如何根據排名獲得其對應的鍵呢?這時候每個節點中的維護的 N 屬性就可以派上用場了。

從根節點向左看,左子樹的節點數就是小於根節點鍵的鍵個數,也就是根節點的鍵排名。所以拿著根節點的左子樹節點數 N 和排名 k 進行比較,會出現三種情況:

  • 左子樹的節點數和排名相等,直接返回根節點的鍵;
  • 左子樹的節點數大於排名,這時候去左子樹接著進行比較;
  • 左子樹的節點數小於排名,說明符合排名要求的節點可能出現在右子樹上(有可能找不到,比如 k 大於整棵二叉樹的節點數),這時候我們得去右子樹搜尋。由於我們直接忽略了左子樹和根節點,所以需要對排名進行一下調整,讓 k = k - N - 1 即可。
@Override
public K select(int k) {
    Node node = root;
    while (node != null) {
        // 父節點左子樹的大小就是父節點的鍵排名
        int N = size(node.left);
        if (N > k) {
            node = node.left;
        } else if (N < k) {
            node = node.right;
            k = k - N - 1;
        } else {
            return node.key;
        }
    }

    return null;
}

根據鍵獲取排名

把根據排名獲取鍵的過程寫作 \(\text{key = select(k)}\),那麼根據鍵獲取排名的過程就是 \(\text{k = select}^{-1}\text{(key) = rank(key)}\)。說明這兩個函式互為反函式。

從根節點出發,拿著根節點的鍵和 key 進行比較會出現三種情況:

  • 根節點的鍵大於 key,這時候得去左子樹中尋找
  • 根節點的鍵小於 key,這時候得去右子樹中尋找,同時得記錄一下左子樹節點數+父節點的那個1
  • 根節點的鍵等於 key,返回根節點的左子樹節點數加上之前跳過的節點數
@Override
public int rank(K key) {
    Node node = root;
    int N = 0;
    while (node != null) {
        int cmp = node.key.compareTo(key);
        if (cmp > 0) {
            node = node.left;
        } else if (cmp < 0) {
            N += size(node.left) + 1;
            node = node.right;
        } else {
            return size(node.left) + N;
        }
    }

    return N;
}

刪除

刪除操作較為複雜,先來看下較為簡單的刪除鍵最小的節點的過程。從根節點出發,一路向左,知道遇到左子樹為 null 的節點,由於這個節點可能還有右子樹,所以需要把右子樹接到父節點上。接完之後還得把這一路上遇到的父節點上的 N - 1。由於沒有其他節點引用了被刪除的節點,所以這個節點會被 java 的垃圾回收機制自動回收。演算法書中給出了一個刪除的示例圖:

刪除最小節點

使用迭代法可以實現尋找最小節點和將右子樹連線到父節點的操作,但是不好處理每一顆子樹的 N 的更新操作,所以還是得靠遞迴法。由於我們需要將最小節點的右子樹接到父節點上,所以滿足終止條件時 deleteMin(Node node) 函式應該把右子樹的根節點返回,否則就應該返回更新之後的節點。

@Override
public void deleteMin() {
    if (root == null) return;
    root = deleteMin(root);
}

private Node deleteMin(Node node) {
    if (node.left == null) return node.right;
    node.left = deleteMin(node.left);
    node.N = size(node.left) + size(node.right) + 1;
    return node;
}

刪除最大的節點的過程和上面相似,只不過我們應該將最大節點的左子樹接到父節點上。

@Override
public void deleteMax() {
    if (root == null) return;
    root = deleteMax(root);
}

private Node deleteMax(Node node) {
    if (node.right == null) return node.left;
    node.right = deleteMax(node.right);
    node.N = size(node.left) + size(node.right) + 1;
    return node;
}

討論完上面兩個較為簡單的刪除操作,我們來看下如何刪除任意節點。從根節點出發,通過比較根節點的鍵和給定的 key,會發生三種情況:

  • 根節點的鍵大於 key,接著去左子樹刪除 key

  • 根節點的鍵小於 key,接著去右子樹刪除 key

  • 根節點的鍵等於 key ,說明我們找到了要被刪除的那個節點,這時候我們又會遇到三種情況:

    • 節點的右子樹為空,直接將左子樹的根節點接到父節點上

    • 節點的左子樹為空,直接將右子樹的根節點接到父節點上

    • 節點的右子樹和左子樹都不為空,這時候需要找到並刪去右子樹的最小鍵節點,然後把這個最小鍵節點頂替即將被刪除節點,把它作為新的子樹根節點

演算法書中給出了第三種情況(右子樹和左子樹都不為空)的示例圖:

刪除

使用遞迴實現的程式碼如下所示:

@Override
public void delete(K key) {
    root = delete(root, key);
}

private Node delete(Node node, K key) {
    if (node == null) return null;

    // 先找到 key 對應的節點
    int cmp = node.key.compareTo(key);
    if (cmp > 0) {
        node.left = delete(node.left, key);
    } else if (cmp < 0) {
        node.right = delete(node.right, key);
    } else {
        if (node.right == null) return node.left;
        if (node.left == null) return node.right;
        Node x = node;
        node = min(x.right);
        // 移除右子樹的最小節點 node,並將該節點作為右子樹的根節點
        node.right = deleteMin(x.right);
        // 設定左子樹的根節點為 node
        node.left = x.left;
    }

    node.N = size(node.left) + size(node.right) + 1;
    return node;
}

總結

如果在插入鍵值對的時候運氣較好,二叉搜尋樹的左右子樹高度相近,那麼插入和查詢的比較次數為 \(\sim2\ln N\) ;如果運氣非常差,差到所有的節點連成了一條單向連結串列,那麼插入和查詢的比較次數就是 \(\sim N\)。所以就有了自平衡二叉樹的出現,不過這已經超出本文的探討範圍了(絕對不是因為寫不動了,以上~~

相關文章