[Python] 트라이 자료구조 구현(Trie)

트라이 자료구조는 검색에 최적화된 자료구조이다

 

자식 노드가 여러 개인 트리 형태로 구성된다고 이해하면 된다.

https://upload.wikimedia.org/wikipedia/commons/thumb/b/be/Trie_example.svg/375px-Trie_example.svg.png

https://leetcode.com/problems/implement-trie-prefix-tree/

Implement Trie (Prefix Tree) - LeetCode

위의 문제를 기준으로 구현했다.

 

class TrieNode:
    
    def __init__(self):
        self.word = False
        self.children = collections.defaultdict(TrieNode)
        
class Trie:

    def __init__(self):
        self.root = TrieNode()

    def insert(self, word: str) -> None:
        node = self.root
        for char in word:
            if char not in node.children:
                node.children[char] = TrieNode()
            node = node.children[char]
        node.word = True
            
    def search(self, word: str) -> bool:
        node = self.root
        for char in word:
            if char not in node.children:
                return False
            node = node.children[char]
        return node.word

    def startsWith(self, prefix: str) -> bool:
        node = self.root
        for char in prefix:
            if char not in node.children:
                return False
            node = node.children[char]
        return True