數(shù)據(jù)結(jié)構(gòu)：前綴樹Trie

引子

在刷題的過程中晌梨，經(jīng)常會遇到這樣一種典型問題：

給一組字符串List<String> strs，找出其中前綴為String p的所有字符串缠沈。

樸素的做法就是遍歷strs维蒙，然后每一個看一下是不是有前綴為p，這樣的效率是O(N*L)版保，其中N是strs中字符串的個數(shù)呜笑，L是p的長度。
這樣看來彻犁，其實也不差叫胁。就單次操作而言，這些操作都是必須的汞幢。問題在于驼鹅，當(dāng)出現(xiàn)多次這種操作的時候，每次都需要重新判斷前綴森篷，這樣效率就不高了谤民。
不賣關(guān)子了，相信大家也知道什么數(shù)據(jù)結(jié)構(gòu)好用了——就是今天要介紹的前綴樹（或者叫字典樹疾宏，又稱單詞查找樹或鍵樹）张足，英文叫做Trie，專門用來處理這種操作坎藐。

介紹

借用百度百科的圖:

Trie

最簡單的Trie的話为牍，只需要保持從當(dāng)前節(jié)點到子節(jié)點的映射就行了，子節(jié)點的字符可以從當(dāng)前節(jié)點得到：

import java.util.HashMap;
import java.util.Map;

class TrieNode {

    private Map<Character, TrieNode> children = new HashMap<Character, TrieNode>();

    public TrieNode() {
    }

    public Map<Character, TrieNode> getChildren() {
        return children;
    }

    public void setChildren(HashMap<Character, TrieNode> children) {
        this.children = children;
    }
}

在此基礎(chǔ)上岩馍，我們可以根據(jù)需要來加入新的變量：

可以用一個Boolean來表示當(dāng)前節(jié)點是不是某個字符串的末尾碉咆，從而在路徑重合的時候知道有哪些字符串；
可以用一個Int來表示當(dāng)前節(jié)點下面完整字符串的個數(shù)蛀恩，用于快速計數(shù)某前綴的個數(shù)疫铜；
可以用一個List來存儲當(dāng)前節(jié)點下面的完整字符串，用于快速返回含有某個前綴的所有字符串双谆。

可以列舉的還有很多壳咕，總之Trie的一大優(yōu)勢就是靈活，按照具體需要進行修改顽馋。

那么以最簡單的Trie為基礎(chǔ)谓厘，如何實現(xiàn)基本的增刪查改操作？

增的話寸谜，比較簡單竟稳，基本上是“順藤摸瓜”，有就一直跟著走，什么地方斷了就自己新建TrieNode他爸。時間復(fù)雜度O(n)聂宾，n是增加的字符串的長度:

    public static void add(String word, TrieNode root) {
        TrieNode cur = root;
        for (char ch : word.toCharArray()) {
            if (cur.getChildren().containsKey(ch)) {
                cur = cur.getChildren().get(ch);
            } else {
                TrieNode node = new TrieNode();
                cur.getChildren().put(ch, node);
                cur = node;
            }
        }
    }

查也不難，思路和增類似诊笤，只是不需要自己創(chuàng)建新節(jié)點而是直接返回false系谐。時間復(fù)雜度還是O(n)：

    public static boolean hasPrefix(String prefix, TrieNode root) {
        TrieNode cur = root;
        for (char ch : prefix.toCharArray()) {
            if (cur.getChildren().containsKey(ch)) {
                cur = cur.getChildren().get(ch);
            } else {
                return false;
            }
        }
        return true;
    }

刪的話相對復(fù)雜一些，一個例子就是兩個完全一樣的字符串盏混，只刪一個怎么實現(xiàn)。上面寫的數(shù)據(jù)結(jié)構(gòu)是不支持重復(fù)字符串的惜论，因此還需要引入一個新變量來存儲對應(yīng)字符串的個數(shù)许赃。這樣就復(fù)雜許多了。
因此馆类，這里的實現(xiàn)假設(shè)字符串都是不重復(fù)的混聊。由于搜索的順序和刪的順序其實是相反的，因此遞歸是不錯的選擇乾巧。刪無非2種情況句喜，需要刪除本節(jié)點，和保留本節(jié)點給其他字符串沟于，所以遞歸還是需要一個返回值的咳胃。復(fù)雜度還是O(n)。

    public static void remove(String word, TrieNode root) {
        removeHelper(word, root, 0);
    }

    private static boolean removeHelper(String word, TrieNode node, int idx) {
        if (idx != word.length()) {
            char ch = word.charAt(idx);
            TrieNode next = node.getChildren().get(ch);
            if (removeHelper(word, next, idx + 1)) {
                node.getChildren().remove(ch);
            }
        }
        return node.getChildren().isEmpty();
    }

改的話旷太，就可以轉(zhuǎn)化成增和刪的組合展懈。

回顧

回到剛開始嘗試解決的問題，假如要找到所有前綴為p的字符串供璧，那么就是遍歷到前綴p的最后一個節(jié)點存崖，之后所有標(biāo)記為完整字符串的就都是了。極限情況下睡毒，前綴p為空来惧，那么相當(dāng)于要遍歷整棵樹，而整棵樹最壞情況就是所有字符串都不重合演顾，那么也就是說復(fù)雜度是O(N*S)供搀，N是字符串個數(shù)，S是字符串的平均長度钠至〕寐可見此時并不優(yōu)越。
那么怎么辦呢棕洋？Trie是可以根據(jù)需要靈活改變的挡闰。在原有Trie的基礎(chǔ)上，執(zhí)行要求的操作效率不高的原因在于威蕉，就算確定了前綴的節(jié)點逞带，后面的字符串還是要一個個去遍歷出來，從而效率不高肃续。很直接的解決思想就是加一個List來維持節(jié)點所對應(yīng)的字符串奢驯。
這樣改變之后申钩，操作的復(fù)雜度就是O(p)，p為前綴p的長度瘪阁，因為只需要找到節(jié)點讀一下List即可撒遣。空間復(fù)雜度的話管跺，其實雖然同樣的字符串在很多節(jié)點出現(xiàn)义黎，但其實都指向同一個對象，其實沒有想象的那么糟糕豁跑×椋或者還可以存字符串List里原來的index。完整代碼如下：

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

class TrieNode {

    private Map<Character, TrieNode> children = new HashMap<Character, TrieNode>();
    private List<String> words = new ArrayList<>();

    public List<String> getWords() {
        return words;
    }

    public void setWords(List<String> words) {
        this.words = words;
    }

    public TrieNode() {
    }

    public Map<Character, TrieNode> getChildren() {
        return children;
    }

    public void setChildren(HashMap<Character, TrieNode> children) {
        this.children = children;
    }

    public static void add(String word, TrieNode root) {
        root.words.add(word);
        TrieNode cur = root;
        for (char ch : word.toCharArray()) {
            if (cur.getChildren().containsKey(ch)) {
                cur = cur.getChildren().get(ch);
            } else {
                TrieNode node = new TrieNode();
                cur.getChildren().put(ch, node);
                cur = node;
            }
            cur.words.add(word);
        }
    }

    public static List<String> getWordsWithPrefix(String prefix, TrieNode root) {
        TrieNode cur = root;
        for (char ch : prefix.toCharArray()) {
            if (cur.getChildren().containsKey(ch)) {
                cur = cur.getChildren().get(ch);
            } else {
                return new ArrayList<>();
            }
        }
        return cur.words;
    }
}

運用

下面看一個例題艇拍，力扣的Word Squares：

值得注意的一點是字符串都是可以重復(fù)使用的狐蜕。

解法與思路

形成Word Squares要求第k行和第k列一致。
第一行的單詞沒有任何限制卸夕。假設(shè)是wall层释。
第二行，k=0快集，那么就是說第0列已經(jīng)確定了湃累，也必須是wall，那么第二行必須以a開始碍讨。假設(shè)是area治力。
第三行，同樣k=0勃黍，k=1都限制了前綴必須是le宵统。假設(shè)是lead。
第四行覆获，限制了前綴必須是lad马澈。
由此可見，每一行都會增加一個限制弄息，前綴必須為第k列的開頭痊班。
因此，一個直接的思路是DFS加back tracking：

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

class Solution {

    public List<List<String>> wordSquares(String[] words) {
        Set<List<String>> ret = new HashSet<>();
        if (words.length == 0) return new ArrayList<>(ret);
        recur(words, new ArrayList<>(), ret);
        return new ArrayList<>(ret);
    }

    private void recur(String[] words, List<String> temp, Set<List<String>> ret) {
        if (temp.size() == words[0].length()) {
            ret.add(new ArrayList<>(temp));
            return;
        }
        for (String w : words) {
            if (meetRequirements(w, temp)) {
                temp.add(w);
                recur(words, temp, ret);
                temp.remove(temp.size() - 1);
            }
        }
    }

    private boolean meetRequirements(String w, List<String> temp) {
        int cnt = temp.size();
        if (cnt == 0) return true;
        for (int i = 0; i < cnt; i++) {
            if (w.charAt(i) != temp.get(i).charAt(cnt)) return false;
        }
        return true;
    }
}

很可惜摹量，這個解法TLE超時涤伐。為何馒胆？因為求下一個單詞時，需要過一遍整個單詞列表來判斷哪些單詞滿足條件凝果，效率不高祝迂。
這時候，我們之前學(xué)的Trie總算派上用場了：可以建一個Trie樹器净，然后需要找前綴的時候在里面搜就好型雳。完整代碼如下(Leetcode AC)：

import java.util.*;

class Solution {

    public List<List<String>> wordSquares(String[] words) {
        Set<List<String>> ret = new HashSet<>();
        if (words.length == 0) return new ArrayList<>(ret);
        TrieNode root = buildTrie(words);
        recur(words, new ArrayList<>(), ret, root);
        return new ArrayList<>(ret);
    }

    private TrieNode buildTrie(String[] words) {
        TrieNode root = new TrieNode();
        for (String s : words) {
            TrieNode.add(s, root);
        }
        return root;
    }

    private void recur(String[] words, List<String> temp, Set<List<String>> ret, TrieNode root) {
        if (temp.size() == words[0].length()) {
            ret.add(new ArrayList<>(temp));
            return;
        }
        // get prefix
        StringBuilder sb = new StringBuilder(temp.size());
        for (int i = 0; i < temp.size(); i++) {
            sb.append(temp.get(i).charAt(temp.size()));
        }
        for (String w : TrieNode.getWordsWithPrefix(sb.toString(), root)) {
            temp.add(w);
            recur(words, temp, ret, root);
            temp.remove(temp.size() - 1);
        }
    }

    private static class TrieNode {

        private Map<Character, TrieNode> children = new HashMap<Character, TrieNode>();
        private final List<String> words = new ArrayList<>();

        public TrieNode() {
        }

        public Map<Character, TrieNode> getChildren() {
            return children;
        }

        public void setChildren(HashMap<Character, TrieNode> children) {
            this.children = children;
        }

        public static void add(String word, TrieNode root) {
            root.words.add(word);
            TrieNode cur = root;
            for (char ch : word.toCharArray()) {
                if (cur.getChildren().containsKey(ch)) {
                    cur = cur.getChildren().get(ch);
                } else {
                    TrieNode node = new TrieNode();
                    cur.getChildren().put(ch, node);
                    cur = node;
                }
                cur.words.add(word);
            }
        }

        public static List<String> getWordsWithPrefix(String prefix, TrieNode root) {
            TrieNode cur = root;
            for (char ch : prefix.toCharArray()) {
                if (cur.getChildren().containsKey(ch)) {
                    cur = cur.getChildren().get(ch);
                } else {
                    return new ArrayList<>();
                }
            }
            return cur.words;
        }
    }
}