字典樹 - Trie | Ian's Blog

什麼是 Trie

在電腦科學中，trie，又稱字首樹或字典樹，是一種有序樹，用於儲存關聯陣列，其中的鍵通常是字串。（取自維基百科）

如同上圖所示， google 關鍵字搜尋的自動填字就是利用 trie 的這個資料結構，利用 trie 的特性，顯示出相關連的字給予選擇。

與 Binary Search Tree、Heap… 等利用 key 來建立的資料結構不同， Trie 最特別的地方就是善用字串的特性：利用每個字的共同前綴(common prefix)當儲存依據。

trie_sample

上圖的 trie 就存了 to, tea, ted, ten, i, in, inn 這些單字。

Basic Trie

現在我們用一個簡單的範例來試試看 trie：

basic_trie = {
    # a and add word
    'a': {
        'd': {
            'd': {'word_end': True},
            'word_end': False},
        'word_end': True},
    # hi word
    'h': {
        'i': {'word_end': True},
        'word_end': False}}


print('Is "a"   a word: {}'.format(basic_trie['a']['word_end']))
print('Is "ad"  a word: {}'.format(basic_trie['a']['d']['word_end']))
print('Is "add" a word: {}'.format(basic_trie['a']['d']['d']['word_end']))

output

Is "a"   a word: True
Is "ad"  a word: False
Is "add" a word: True

範例中有一個基本的字典，結構類似於上面那張圖，裡面有三個單字分別是 a、add以及hi。當要查一個單字時，會依序把每個字母傳進去，如果最後傳word_end沒有字母了。可以發現當要搜尋 ad 時，因為這在 trie 內並不算一個單字，所以最後回傳 False。

我們可以改良一下上述的程式碼

def is_word(word):
    current_node = basic_trie
    
    for char in word:
        if char not in current_node:
            return False
        
        current_node = current_node[char]
    
    return current_node['word_end']


# Test words
test_words = ['ap', 'add']
for word in test_words:
    if is_word(word):
        print('"{}" is a word.'.format(word))
    else:
        print('"{}" is not a word.'.format(word))

藉由迴圈來判斷這個單字的每個字母是否在這個 trie 內，如果沒有字母就取 word_end 來判斷他是不是一個單字。

實作 Trie

在了解 trie 的資料結構之後，現在來把它實作出來。

由於 trie 需要用 node 來建立，所以要先把 node 的 class 給建立起來。剛剛有提到，每個 node 裡面存放一個前贅字，所以需要一個 children 的 dict 物件。接著若現在不是為最後一個字母，那就不是一個完整的單字，所以需要一個 is_word 的布林值。

class TrieNode(object):
    def __init__(self):
        self.is_word = False
        self.children = {}

有了 node 後，接著就要建立 Trie 的物件了，我們預期這個字典有兩個功能，一個是新增單字（add）的功能，另一個則是查詢單字是否存在於此字典（exists）的功能。

add(word)

依照剛剛了解的資料結構，必須要將單字拆成逐一的字母放入 children 內形成樹狀結構，並在最後把 is_word 設定成 True。

exists(word)

要檢查單字是否存在於字典內，也是同樣的奧里，將字母逐一比對 children 的值，若存在則往下，不存在則取回 is_word。

class Trie(object):
    def __init__(self):
        self.root = TrieNode()

    def add(self, word):
        """
        Add `word` to trie
        """
        current_node = self.root

        for char in word:
            if char not in current_node.children:
                current_node.children[char] = TrieNode()
            current_node = current_node.children[char]

        current_node.is_word = True

    def exists(self, word):
        """
        Check if word exists in trie
        """
        current_node = self.root

        for char in word:
            if char not in current_node.children:
                return False
            current_node = current_node.children[char]

        return current_node.is_word

最後就是用下面的 test case 來測試程式是否正確。


# -- Test Case --

word_list = ['apple', 'bear', 'goo', 'good', 'goodbye', 'goods', 'goodwill', 'gooses'  ,'zebra']
word_trie = Trie()

# Add words
for word in word_list:
    word_trie.add(word)

# Test words
test_words = ['bear', 'goo', 'good', 'goos']
for word in test_words:
    if word_trie.exists(word):
        print('"{}" is a word.'.format(word))
    else:
        print('"{}" is not a word.'.format(word))

output

"bear" is a word.
"goo" is a word.
"good" is a word.
"goos" is not a word.

這樣就簡單實作了一個 trie 了。

練習程式碼

可以到 clone github 的程式碼練習（06），或是把程式碼複製到 online_python_compiler 練習也可以。

Resource

Trie