今天這篇文章源於上周在工作中解決的一個實際問題，它是個比較普遍的問題，無論做什么開發，估計都有遇到過。具體是這樣的，我們有一份高校的名單（2657個），需要從海量的文章標題中找到包含這些高校的標題，其實就是模糊查詢（關注公眾號 渡碼， 回復關鍵詞 trie 獲取源碼）。對應的偽代碼如下
selected_titles = [] for 標題 in 海量標題： for 高校 in 高校名單： if 標題.contains(高校): selected_titles.add(標題) break
如果是大數據開發，對應的SQL的偽代碼是這樣的
select title from tb where title rlike '清華大學|北京大學|...2657個高校'
上面這兩種做法都能實現我們的需求，但它們的共同問題是查詢效率太低。如果我們要匹配的高校不是2657個而是幾十萬甚至上百萬，那這種方式耗費時之久是不可想象的。
優化這類問題通常需要在數據結構上做文章，這個問題中我們能優化的數據結構也只有“高校名單”這個了，上面的偽代碼中我們存放“高校名單”的數據結構是數組，當我們查找某個title是否包含某個高校的時候，需要從頭到尾遍歷一遍“高校名單”，並且名單越長，遍歷耗時就越長。
清楚了數組這種數據結構的缺點后，接下來我們重點要做的就是尋找一個數據結構可以做到在不遍歷整個“高校名單”的情況下就可以完成模糊查詢。這個數據結構就是我們今天要介紹的 Trie 樹，冷眼一看這個單詞有點陌生，又是一個樹型結構，感覺會很復雜似的，實際上這個數據結構的設計思想非常簡單，一學就會。
下面我們就來學習一下 Trie 樹。為了方面講解，假設“高校名單”里只有下面5個元素
ABC、ABD、BCD、BCE、C、CAB、CDE
對應的兩種數據結構如下：
拋開這兩種數據結構查找的時間復雜度，我們先從直觀上看看為什么 Trie 樹的查找效率要比數組高。假設我們要查找，“CDE”這個字符串，在數組結構中，我們要遍歷一遍數組，比較7次才能找到結果，做了比較多的“無用功”。而在 Tire 樹中只需要比較3次就可以找到，它的優勢非常明顯，由於樹型結構我們根本不用考慮左側A、B開頭的兩個分支，這就大大減少了比較的次數，從而減少“無用功”。下面用一個動畫來演示一下如何創建 Trie 樹，以及在 Trie樹上查找字符串（如果視頻播放不了可以看源碼目錄中的gif）
樹的建立過程其實就是遍歷字符串每個元素並在樹上建立相應的節點。字符串查找過程其實就是按照字符串對樹進行遍歷。Trie 樹的建立與字符串查找還是比較簡單的。
不知道大家是否注意到上圖 Trie 樹中的節點有兩種顏色——白色和綠色。綠色節點代表從根節點到當前節點的字符串是“高校名單”中的字符串，也就是我們建立 Trie 樹用到的字符串。以最左側的葉子結點“C”為例，它代表“ABC”字符串是“高校名單”中的字符串。同理，字符串“AB”就不是“高校名單”里的元素，因為“B”節點不是綠色的，因此當我們在這棵樹上查找字符串“AB”時，是查不到的。這一點需要大家注意，下面編碼中我們也有體現。
另外，有朋友可能會有疑問，我們最開始的需求不是模糊查詢嗎，在 Trie 樹講解這部分怎么都在說字符串全詞（精確）匹配。這是因為全詞匹配是 Tire 樹支持的最基本的查找方式，在此基礎上，我們做一些變通就可以很容易實現模糊匹配。
接下來，我們就來看看代碼實現（Python版），首先創建兩個數組
colleges = utils.read_file_to_list('key_words.txt') titles = utils.read_file_to_list('titles.txt')
colleges就是我們一直在說的“高校名單”，titles便是“海量標題”，它們都是一維數組，數組每個元素都是一個字符串。
再來編寫 Trie 樹相關的代碼，如果理解了 Trie 樹的設計思想，再編寫下面的代碼其實很容易。首先要定義一個類代表 Trie 樹節點
class TrieNode: def __init__(self): self.nodes = dict() # is_end=True 代表從根節點到當前節點構造Trie樹的字符串（出現在“高校名單”里）
self.is_end = False
is_end=True就是我們上面說的綠色節點。
再來編寫創建 Trie 樹的代碼，代碼在 TrieNode 類中
def insert_many(self, items: [str]): """ 支持輸入字符串數組，直接構造一個 Trie 樹 :param items: 字符串數組 :return: None """
for word in items: self.insert(word) def insert(self, item: str): """ 向 Trie 樹插入一個短語 :param item: 待插入的字符串 :return: None """ curr = self for word in item: if word not in curr.nodes: curr.nodes[word] = TrieNode() curr = curr.nodes[word] curr.is_end = True
再來編寫查找 Trie 樹的代碼，代碼在 TrieNode 類中
def suffix(self, item: str) -> bool: """ 匹配前綴，也就是判斷item字符串是否是以“高校名單”中某個字符串開頭 :param item: 待匹配字符串 :return: True or False """ curr = self for word in item: if word not in curr.nodes: return False curr = curr.nodes[word] # 取得子節點
if curr.is_end: # 如果is_end=True說明當前字符串包含了“高校名單”的某個字符串
return True return False # 未匹配上
這里並不是全詞匹配，而是前綴匹配，也就是判斷待查找的字符串item是否是以“高校名單”中某個字符串開頭。
再來編寫模糊匹配，代碼在 TrieNode 中
def infix(self, item: str) -> bool: for i in range(len(item)): sub_item = item[i:] # 將待查找的字符串分成不同子串
# 如果子串的前綴在 Trie 樹中能匹配上
# 說明待查找的字符串item中包含“高校名單”中的元素，
# 即實現了 tile rlike '清華大學|北京大學|...其他大學' 的功能
if self.suffix(sub_item): return True return False
這里其實就是把待查找字符串item分成不同子串去做前綴匹配，如果子串匹配上，那就說整個字符串item就包含了“高校名單”里面的某個字符串。
最后，我們運行一下上面的代碼，並記錄查找時間，與最開始數組結構那一版做個對比。代碼如下
# 數組版本
cnt = 0 start_time = int(time.time() * 1000) for title in titles: for x in colleges: if x in title: cnt += 1
break end_time = int(time.time() * 1000) print(cnt) print('spend: %.2fs' % ((end_time - start_time) / 60.0)) # Trie 樹版本
root = TrieNode() root.insert_many(colleges) cnt = 0 start_time = int(time.time() * 1000) for title in titles: if root.infix(title): cnt += 1 end_time = int(time.time() * 1000) print(cnt) print('spend: %.2fs' % ((end_time - start_time) / 60.0))
輸出結果如下：
5314 spend: 9.13s 5314 spend: 0.23s
可以看到，用數組匹配用了9s，而用 Trie 樹匹配僅用0.23s！
今天介紹的這種提高海量數據模糊查詢性能的方式是通過寫代碼的方式實現的，對於經常寫 SQL 的大數據開發者來說，要把它用起來只是建個 UDF 就可以了，需要在 UDF 的初始化代碼中用“高校名單”建立一顆 Trie 樹。
今天的內容就分享到這里了，希望對你有幫助。公眾號回復關鍵詞 trie 獲取完整源代碼
歡迎公眾號「渡碼」，輸出別地兒看不到的干貨。
