트라이(trie) 파이썬

https://qkqhxla1.tistory.com/700

그림을 다시 가져옴.

아주 예전에 정리했었는데 c++이고.. 좀 복잡해서 지나면 잊어버린다. 파이썬 딕셔너리로 쉽게 구현한 버전이 있어서 가져왔다. MyPrettyPrinter클래스는 디버깅용이니 무시하고 Trie클래스만 보자.

소스 출처 : https://leetcode.com/problems/add-and-search-word-data-structure-design/discuss/59700/Python-trie-solution-search-using-dfs 여기서 가져와서 보기좋게 가공했다.

addWord로 단어가 추가되면 트라이를 구성하는데, 단어가 어떻게 구성되는지는 아래에

print MyPrettyPrinter().pformat(t.trie) 로 출력하게 해놔서(글 맨 아래에도 출력 결과 있음.) 눈으로 대충 확인이 가능하다.

# -*- coding: utf-8 -*-

import pprint
class MyPrettyPrinter(pprint.PrettyPrinter):
    def format(self, _object, context, maxlevels, level):
        if isinstance(_object, unicode):
            return "'%s'" % _object.encode('utf8'), True, False
        elif isinstance(_object, str):
            _object = unicode(_object,'utf8')
            return "'%s'" % _object.encode('utf8'), True, False
        return pprint.PrettyPrinter.format(self, _object, context, maxlevels, level)
    

class Trie:
    def __init__(self):
        self.trie = {}

    def addWord(self, word):
        trie = self.trie
        for c in word:
            if c not in trie:
                trie[c]={}
            trie = trie[c]
        trie['#'] = '#'  # 끝을 나타내는 #.

    def search(self, word):  # 탐색.
        stack = [[word, self.trie]]
        while stack:
            word, trie = stack.pop()
            if not word:  # word를 다 돌았는데
                if '#' in trie:  # 끝이면
                    return True  # return True
                continue  # 아니면 다른것도 돌려봄.

            c = word[0]
            if c in trie:  # 그 다음 단어가 trie안에 있으면 그 다음단어를 찾으러 감.
                stack.append([word[1:], trie[c]])
        return False

t = Trie()
for word in ["ran", "rune", "runner", "runs", "add", "adds", "adder", "addee"]:
    t.addWord(word)
print MyPrettyPrinter().pformat(t.trie)  # trie확인용. 디버깅용.
print t.search('runner')

중간에 디버깅용으로 출력한 trie의 모습은 아래와 같다.

{'a': {'d': {'d': {'#': '#',

                   'e': {'e': {'#': '#'}, 'r': {'#': '#'}},

                   's': {'#': '#'}}}},

 'r': {'a': {'n': {'#': '#'}},

       'u': {'n': {'e': {'#': '#'},

                   'n': {'e': {'r': {'#': '#'}}},

                   's': {'#': '#'}}}}}

딕셔너리안에 딕셔너리가 있는 형태인데 모든 단어가 a와 r로 시작한다. 그리고 그 다음에 ad로 시작하는거, ra, ru로 시작하는거.. 이렇게 계속 가다가 #이 끝이라는 소리이다. {'a': {'d': {'d': {'#': '#',의 경우에는 add라는 단어를 나타낸다.

t = Trie()
for word in ["ran", "rune", 'ade']:
    t.addWord(word)
print MyPrettyPrinter().pformat(t.trie)  # trie확인용.
print t.search('ran')

로 바꿔서 돌려보면 

{'a': {'d': {'e': {'#': '#'}}},

 'r': {'a': {'n': {'#': '#'}}, 'u': {'n': {'e': {'#': '#'}}}}}

가 출력된다. 너무 이해하기가 쉬워서 붙여놓는다.