트라이(TRIE)

[Trie란]

Trie(retrieval: 검색)는 검색 트리의 일종으로 일반적으로 문자열의 각 문자를 키로 하여 순차적으로 문자열의 구성 정보를 저장하는 데 사용한다. 주로 검색어 자동완성, 사전 찾기, 문자열 검사 등에 사용된다.

 

---

[Trie의 시간복잡도]

• insert: 문자열 개수를 N라 하고 가장 긴 문자열의 길이를 L이라 할 때 Trie에 모든 데이터를 삽입하는 작업의 시간 복잡도는  O(N * L)이다.
• search: 찾고자 하는 문자열의 길이가 L이라고 할 때 해당 문자열의 유무를 확인하는 시간 복잡도는 O(L)이다.

 

---

 

[Trie의 구현]

from collections import defaultdict


class TrieNode:
    def __init__(self):
        self.word = ''
        # defaultdict() 기능 확인.
        self.children = defaultdict(TrieNode)


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

    def insert(self, string: str) -> None:
        node = self.root
        for char in string:
            node = node.children[char]
        node.word = string

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

        return True if node.word else False

    def starts_with(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

    def auto_complete(self, prefix: str) -> list:
        node = self.root
        candidates = []
        for char in prefix:
            if char not in node.children:
                return candidates
            node = node.children[char]

        if node.word:
            candidates.append(prefix)

        stack = list(node.children.values())

        while stack:
            v = stack.pop()

            if v.word:
                candidates.append(v.word)

            stack.extend(list(v.children.values()))
        return candidates


my_trie = Trie()
my_trie.insert('sarang')
my_trie.insert('saram')
my_trie.insert('satang')
my_trie.insert('sara')

print(my_trie.search('sa'))         # 일치하는 문자열이 없으므로 False
print(my_trie.search('sarang'))     # 일치하는 문자열이 존재하여 True

print(my_trie.starts_with('ta'))    # 'ta'로 시작하는 문자열이 없으므로 False
print(my_trie.starts_with('sa'))    # 'sa'로 시작하는 문자열이 존재하므로 True

print(my_trie.auto_complete(''))    # 아무 것도 입력되지 않으면 모든 문자열을 보여줌.
print(my_trie.auto_complete('sar')) # 입력 문자열이 'sar'이면 sar로 시작하는 문자열들을 보여 줌.

 

위 코드르 실행한 결과는 다음과 같다.

False
True
False
True
['satang', 'sara', 'saram', 'sarang']
['sara', 'saram', 'sarang']