[Python] Trie 자료구조

Trie(트라이) 자료구조는 문자열 집합에서 키를 효율적으로 저장하고 조회하는데 사용된다.

스펠 체커나 자동완성 기능에서 사용된다.

파이썬에서는 딕셔너리 자료구조를 이용해 쉽게 구현할 수 있다.

• insert (단어 삽입)
  • 단어의 키(문자)가 커서에 존재하지 않는 경우 해당 키로 딕셔너리를 생성하고, 커서 위치를 업데이트 해준다.
  • 단어 삽입이 완료되면 * 키를 생성해서 끝임을 표기한다.
• search (단어 조회)
  • 단어의 키(문자)가 커서에 존재하지 않는 경우 False를 반환하고, 커서 위치를 업데이트 해준다.
  • 단어 순회가 완료되었을 때 * 키가 커서에 있을 경우 True를 반환하고 없을 경우 False를 반환한다.
• startswith (접두어 조회)
  • 단어의 키(문자)가 커서에 존재하지 않는 경우 False를 반환하고, 커서 위치를 업데이트 해준다.
  • 단어 순회가 완료되었을 때 True를 반환한다.

class Trie:
    def __init__(self):
        self.trie = dict()

    def insert(self, word: str) -> None:
        cur = self.trie
        for ch in word:
            if ch not in cur.keys():
                cur[ch] = dict()
            cur = cur[ch]
        cur['*'] = True

    def search(self, word: str) -> bool:
        cur = self.trie
        for ch in word:
            if ch not in cur.keys():
                return False
            cur = cur[ch]
        if '*' in cur.keys():
            return True
        else:
            return False

    def startsWith(self, prefix: str) -> bool:
        cur = self.trie
        for ch in prefix:
            if ch not in cur.keys():
                return False
            cur = cur[ch]
        return True