今天這篇文章源於上週在工作中解決的一個實際問題，它是個比較普遍的問題，無論做什麼開發，估計都有遇到過。具體是這樣的，我們有一份高校的名單（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 獲取完整原始碼

歡迎公眾號「渡碼」，輸出別地兒看不到的乾貨。

&n