數據結構之字典樹 Trie 單詞查找樹

參照算法第4版，強烈建議看書上的實現步驟，這裏只是一個個人記錄。

單詞查找樹的性能：

1. 查找命中所需的時間與被查找的鍵的長度成正比。比如單詞有7個字符，查找或插入操作最多隻需要檢查8個節點。
2. 查找未命中只需檢查若干個字符。
    

所需空間：在RN~RNw之間，其中R爲字母表大小，N爲鍵的個數，w爲平均單詞長度。

import java.util.LinkedList;
import java.util.Queue;

/**
 * @author yuan
 * @date 2019/2/21
 * @description 單詞查找樹，參考算法第4版
 */
public class TrieST<Value> {
    /**
     * 基數
     */
    private static int R = 256;

    /**
     * 根節點
     */
    private Node root;

    private static class Node{
        private Object val;
        private Node[] next = new Node[R];
    }

    public Value get(String key) {
        Node x = get(root, key, 0);
        if (x == null) {
            return null;
        }
        return (Value) x.val;
    }

    /**
     * 返回以x作爲根節點的子單詞查找樹中與key相關聯的值
     * @param x
     * @param key
     * @param d 當前的遍歷深度
     * @return
     */
    private Node get(Node x, String key, int d) {
        if (x == null) {
            return null;
        }
        if (d == key.length()) {
            return x;
        }
        char c = key.charAt(d);
        return get(x.next[c], key, d + 1);
    }

    public void put(String key,Value val) {
        root = put(root, key, val, 0);
    }

    /**
     * 如果key存在以x爲根節點的子單詞查找樹中則更新與它相關聯的值
     * @param x
     * @param key
     * @param val
     * @param d
     * @return
     */
    private Node put(Node x, String key, Value val, int d) {
        if (x == null) {
            x = new Node();
        }
        if (d == key.length()) {
            x.val = val;
            return x;
        }
        // 找到第d個字符所對應的子單詞查找樹
        char c = key.charAt(d);
        x.next[c] = put(x.next[c], key, val, d + 1);
        return x;
    }

    /**
     * 查找所有的鍵
     * @return
     */
    public Iterable<String> keys(){
        return keysWithPrefix("");
    }

    /**
     * 匹配以pre爲前綴的鍵
     * @param pre
     * @return
     */
    private Iterable<String> keysWithPrefix(String pre) {
        Queue<String> q = new LinkedList<>();
        collect(get(root, pre, 0), pre, q);
        return q;
    }

    private void collect(Node x, String pre, Queue<String> q) {
        if (x == null) {
            return;
        }
        if (x.val != null) {
            q.offer(pre);
        }
        for (char c = 0; c < R; c++) {
            collect(x.next[c], pre + c, q);
        }
    }

    /**
     * 匹配含通配符的鍵('.'能匹配所有字符)
     * @param pat
     * @return
     */
    public Iterable<String> keysThatMatch(String pat) {
        Queue<String> q = new LinkedList<>();
        collect(root, "", pat, q);
        return q;
    }

    private void collect(Node x, String pre, String pat, Queue<String> q) {
        int d = pre.length();
        if (x == null) {
            return;
        }
        if (d == pat.length() && x.val != null) {
            q.offer(pre);
        }
        if (d == pat.length()) {
            return;
        }
        char next = pat.charAt(d);
        for (char c = 0; c < R; c++) {
            if (next == '.' || next == c) {
                collect(x.next[c], pre + c, pat, q);
            }
        }
    }

    /**
     * 查找s中的最長的鍵
     *
     * @param s
     * @return
     */
    public String longestPrefixOf(String s) {
        int length = search(root, s, 0, 0);
        return s.substring(0, length);
    }

    private int search(Node x, String s, int d, int length) {
        if (x == null) {
            return length;
        }
        // 當前節點非空，更新length
        if (x.val != null) {
            length = d;
        }
        // 已經遍歷到s的末尾，返回
        if (d == s.length()) {
            return length;
        }
        char c = s.charAt(d);
        return search(x.next[c], s, d + 1, length);
    }

    /**
     * 刪除操作
     * @param key
     */
    public void delete(String key) {
        root = delete(root, key, 0);
    }

    private Node delete(Node x, String key, int d) {
        if (x == null) {
            return null;
        }
        if (d == key.length()) {
            x.val = null;
        } else {
            char c = key.charAt(d);
            x.next[c] = delete(x.next[c], key, d + 1);
        }
        if (x.val != null) {
            return x;
        }
        for (char c = 0; c < R; c++) {
            if (x.next[c] != null) {
                return x;
            }
        }
        return null;
    }

    public static void main(String[] args) {
        TrieST<Integer> trieST = new TrieST<>();
        trieST.put("abc", 3);
        trieST.put("acc", 1);
        trieST.put("bb", 1);

        System.out.println(trieST.get("abc")); // 3

        System.out.println();

        Iterable iterable = trieST.keysThatMatch("a.c");
        iterable.forEach(i -> System.out.println(i)); // abc acc

        System.out.println();

        trieST.delete("abc");
        System.out.println(trieST.get("abc")); // null

    }
}