[파이썬] Trie(트라이) 알고리즘

참고한 블로그

[Python / 파이썬] Trie 알고리즘

 

Trie 알고리즘 

 -- 문자열 탐색 시 실행속도가 빠른 알고리즘. 하지만 저장 공간의 크기가 크다는 단점.

 

---

구현

 

데이터를 저장할 노드

class Node(object):
    def __init__(self,key,data = None):
        self.key = key
        self.data = data
        self.children = {}

 

Trie 구현

class Trie:
    def __init__(self):
        self.head = Node(None)

    def insert(self,string):
        current_node = self.head
        for char in string:
            if char not in current_node.children:
                current_node.children[char] = Node(char)
            current_node = current_node.children[char]    
        current_node.data = string
        
    def search(self,string):
        current_node = self.head
        for char in string:
            if char not in current_node.children:
                return False
            current_node = current_node.children[char]
        if current_node.data:
            return True
        return False
    
    def starts_with(self,prefix):
        current_node = self.head
        words = []
        for p in prefix:
            if p not in current_node.children:
                return False
            current_node = current_node.children[p]
        
        current_node = [current_node]    
        next_node = []
        
        while True:
            for node in current_node:
                if node.data:
                    words.append(node.data)
                next_node.extend(list(node.children.values()))
            if len(next_node) != 0:
                current_node = next_node
                next_node = []
            else:
                break
        return words

 

실행 예시

 

###example###    
trie = Trie()
words = ["frodo","firefox","fire","friedchicken"]

for word in words:
    trie.insert(word)

print(trie.search("frodo"))
print(trie.starts_with("f"))

 

 

실행 결과

True
['fire', 'frodo', 'firefox', 'friedchicken']