[leetcode] 208. Implement Trie (Prefix Tree) 파이썬 코드

prefix를 담은 tree를 trie ('트라이')라고 한다. 자동 완성이나 스펠링 체크를 할 때 쓰이는 알고리즘이다. 

트라이를 구현해보자.

 

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

class Trie(object):

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

    def insert(self, word):
        cur = self.root

        for c in word:
            if c not in cur.children: # 만약 c가 자식 노드에 없다면
                cur.children[c] = TrieNode() # 추가해준다
            cur = cur.children[c] # 그리고 자식 노드로 넘어간다

        cur.endOfWord = True # 마지막 단어라는 표시를 해줘야 한다

    def search(self, word):
        cur = self.root

        for c in word:
            if c not in cur.children: # 자식노드에 없다면
                return False # false 리턴
            cur = cur.children[c] # 그 다음 단어로 넘어가기

        return cur.endOfWord # 그것이 마지막 글자였는지 확인해야 함
        

    def startsWith(self, prefix):
        cur = self.root

        for c in prefix:
            if c not in cur.children:
                return False
            cur = cur.children[c]
        
        return True # search와는 다르게 마지막 단어인지 확인할 필요없이 그것이 포함된 단어가 있으면 True 리턴

 

 

참고

https://www.youtube.com/watch?v=oobqoCJlHA0&t=601s