數據結構 | 30行代碼，手把手帶你實現Trie樹

時間 2020-07-20

標籤數據結構代碼手把手實現 trie 简体版

原文原文鏈接

本文始發於我的公衆號：TechFlow，原創不易，求個關注node

今天是算法和數據結構專題的第28篇文章，咱們一塊兒來聊聊一個經典的字符串處理數據結構——Trie。web

在以前的4篇文章當中咱們介紹了關於博弈論的一些算法，其中應用最廣也是最重要的就是最後的SG函數。瞭解到這些以後，足夠咱們應付常見的博弈論算法問題了。博弈論自己就是一門學科，其中有這很深邃的理論基礎，咱們只是淺嘗輒止，你們感興趣的能夠自行鑽研一下，相信必定會頗有收穫。面試

小故事

之前讀過一個大牛的文章，文章裏討論了一個問題，若是不是爲了面試的話，咱們爲何要學算法？算法

他講了一個他本身的故事，說是在不少年前，手機仍是諾基亞功能機的時代，他爲塞班系統開發了一個通信簿查找聯繫人的軟件。軟件的功能很簡單，就是存儲聯繫人，而後能夠經過拼音或者是拼音首字母查找到對應的聯繫人。這裏須要對漢字以及拼音的映射作一個處理，也不是很複雜的操做，咱們腦補應該就能夠想出來。數組

軟件很快作好了，作好了以後投入使用發現也很好用。可是很快遇到了一個沒想到的問題，就是當聯繫人多了以後，軟件的運行速度變得很是慢，也就是卡。卡的緣由也很簡單，由於搜索聯繫人的這個步驟他用的是遍歷查找的方式搜索的。他一開始先是本身腦補了一些優化方案和野路子，雖然能有些提高可是不能根本解決問題。後來被逼無奈，他在搜索了相關資料以後，找到了咱們今天的主角Trie，用上了這個算法以後，這個問題瞬間迎刃而解，即便存儲了成千上萬的聯繫人不再會卡頓了。今後他大徹大悟，算法並非奇淫技巧，真的是有用的。數據結構

咱們就以這個文章當中的問題做爲基礎，來看看Trie的原理，以及它爲何能夠解決這個問題。機器學習

Trie簡介

Trie樹有好幾個中文翻譯名字，有的稱爲字典樹，有的稱爲前綴樹。這都是能夠的，看你們各位的喜歡。我通常就稱Trie樹，對方聽不懂纔會說字典樹XD。編輯器

從字典樹和前綴樹的稱謂當中咱們是能夠腦補出來它的大概原理的，也就是以字典和前綴的形式存儲數據。聽着有點抽象，其實咱們看一張圖就徹底明白了。函數

從圖中咱們能夠看出來，Trie樹是一棵多叉樹。樹中的每個節點存儲一個字符，咱們從根節點到節點的路徑上的字符連起來就成了單詞。也就是說全部的單詞都是這樣縱向的形式存儲在樹上的。學習

這樣存儲有什麼好處呢？最大的好處就是擁有相同前綴的單詞能夠共享前綴，好比ana，ann和and這兩個單詞前兩個字符是相同的都是an，因此他們擁有一條公共前綴鏈路。在這樣的結構之下，當咱們須要查找單詞是否在樹中的時候，咱們只須要從樹根開始遍歷，若是能找到對應的結尾節點就說明單詞存在，不然就說明單詞不存在。

舉個例子，好比咱們要查找單詞doe，咱們先從根節點查找字符d。發現d存在，因而轉移到了d。接着咱們查找字符o，在d節點下面也找到了字符o，轉移到o，查找字符e，發現e不存在，因而就說明了doe單詞不存在。可是這裏有一個問題，假設咱們存在一個單詞是doea，咱們查找doe仍是能夠找到，可是doe單詞實際上是不存在的，這不就錯誤了嗎？

的確，這樣可能存在問題，因此咱們須要在節點當中記錄一下，是不是某一個單詞的結尾。這樣咱們不只須要找到對應的單詞，還須要防止咱們將其餘單詞的前綴當成是單詞。

咱們插入單詞的過程和查詢很是接近，一樣是一個樹上遍歷的過程，只不過若是咱們發現查詢的節點不存在時會手動建立。整個單詞插入完成以後，將最後一個字符對應的節點進行標記，代表這是一個單詞的結尾。

簡單的Trie樹只須要完成添加和查詢便可，若是要涉及刪除，咱們只須要在節點當中維護一下通過該節點的單詞數量便可。在刪除的時候，將沿途通過的節點標記的數量-1，若是遇到數量爲0的節點，直接刪除便可。

代碼實現

光說不練假把式，咱們天然也是要來練練的。

相信你們也從描述當中看出來了，Trie的原理說穿了其實很簡單，實現起來也不困難。網上有許多版本，不少是面向過程實現的，我把它封裝了一下，用Python面向對象實現了一個版本。理解了原理以後，你們能夠根據本身的須要開發本身的風格的版本，代碼其實不過重要，主要仍是理解原理。

我把Trie樹分紅了兩個部分，第一個部分是樹上的節點。對於Trie樹上的節點來講它須要提供兩個功能。第一個功能是返回當前節點是不是某一個字符串的結尾，第二個功能是根據字符查找後繼的節點。咱們只須要在類當中設置一個flag標記和一個dict屬性來存儲後繼元素就好了。

class Node:
    def __init__(self, is_leaf=False):
        self.child = {}
        self._is_leaf = is_leaf

 @property
    def is_leaf(self):
        return self._is_leaf

 @is_leaf.setter
    def is_leaf(self, is_leaf):
        self._is_leaf = is_leaf

    # 加入孩子節點
    def put(self, key, value):
        self.child[key] = value

 @staticmethod
    # 將字符轉化成數字
    # 其實沒有必要，由於用到了dict，若是用數組存儲孩子的話，須要用它來計算下標
    def get_idx_of_str(_str):
        if len(_str):
            return -1
        if ord('a') <= ord(_str) <= ord('z'):
            return ord(_str) - ord('a')
        else:
            return ord(_str) - ord('A') + 26

    # 根據字符獲取下一個節點
    def get_next_node(self, _str):
        if len(_str) != 1:
            return None
        idx = Node.get_idx_of_str(_str)
        return self.child.get(idx, None)

    def get_node(self, key):
        return self.child.get(key, None)

這裏我將is_leaf兩個方法用property封裝，從而能夠方便使用，這個也是經常使用的慣例。有了節點以後，咱們再開發Trie類就很方便了，對於Trie這個類而言咱們只須要實現兩個方法，一個是插入字符串，一個是字符串的查詢。在有了Node類以後，這兩個方法實現也很簡單了。

class Trie:
    def __init__(self):
        self.root = Node()

    def insert(self, _str):
        cur = self.root
        # 遍歷字符
        for c in _str:
            # 查找下一個節點
            if cur.get_next_node(c) is None:
                # 若是節點不存在，本身建立一個新節點並插入
                key = Node.get_idx_of_str(c)
                cur.put(key, Node())
                cur = cur.get_node(key)
            # 不然繼續往下
            else:
                cur = cur.get_next_node(c)
        cur.is_leaf = True
    
    def query(self, _str):
        cur = self.root
        # 遍歷字符
        for c in _str:
            # 查詢，若是查詢不到返回False
            if cur.get_next_node(c) is None:
                return False
            cur = cur.get_next_node(c)
        # 返回是不是字符串結尾
        return cur.is_leaf

這兩段代碼應該都不能讀懂，最後，咱們嘗試一下使用它來測試一下：

if __name__ == "__main__":
    trie = Trie()
    trie.insert('abcda')
    trie.insert('abcde')
    trie.insert('eecdab')
    trie.insert('mout')
    trie.insert('ymm')
    print(trie.query('abcda'))
    print(trie.query('mout'))
    print(trie.query('ym'))