[Algorithm] trie 자료구조 python

알고리즘 문제를 풀다가 trie 자료구조가 필요한 문제를 접했다.

2020 카카오 신입 공채 1차
https://programmers.co.kr/learn/courses/30/lessons/60060

fro?? 가 있으면 words 중 다섯 글자인 단어들만 뽑고, 그 중에서 맨처음부터 fro를 포함하는 단어를 찾도록 구현했는데, 정확성에서는 통과를 했지만 효율성에서 통과를 하지 못했다.

좀 더 고민을 했어야 하긴 하지만 어떻게 풀어야할지 궁금해서 찾아보니 trie라는 자료구조를 활용해야 된다고 한다.

trie

halo가 존재하는지 확인하려면

1. head 노드의 자식 노드 중에 h가 있는 것을 확인한 뒤 h로 이동
2. h의 자식 노드 중에 a가 있는 것을 확인하고 a로 이동
3. a의 자식 노드 중에 l이 있는 것을 확인하고 l로 이동
4. l의 자식 노드 중 o가 있는 것을 확인하고 o로 이동
5. o의 자식 노드가 없다는 것을 확인

이 과정을 거쳐 trie에 halo가 있다는 것을 확인할 수 있다.
o의 자식 노드가 없음을 표시하기 위해 o의 자식 노드로 * 를 생성하면 편하다.

python 코드

class Trie:
    head = {}

    def add(self, word):
        cur = self.head
        for ch in word:
            if ch not in cur:
                cur[ch] = {}
            cur = cur[ch]
        cur['*'] = True

    def search(self, word):
        cur = self.head
        for ch in word:
            if ch not in cur:
                return False
            cur = cur[ch]
        if '*' in cur:
            return True
        else:
            return False

dictionary = Trie()

dictionary.add("frodo")
dictionary.add("front")
dictionary.add("frost")
dictionary.add("frozen")
dictionary.add("frame")
dictionary.add("kakao")

print(dictionary.search("frodo"))
print(dictionary.search("froze"))
print(dictionary.search("kaka"))

# 생성 결과

{'f': {'r': {'o': {'d': {'o': {'*': True}}, 'n': {'t': {'*': True}}, 's': {'t': {'*': True}}, 'z': {'e': {'n': {'*': True}}}}, 'a': {'m': {'e': {'*': True}}}}},
'k': {'a': {'k': {'a': {'o': {'*': True}}}}}}

코드를 보면 바로 이해할 수 있을 것이다.

add 함수에서는 word의 알파벳 하나하나씩 추가해주고, 마지막에 * 노드를 추가해준다.

search 함수에서는 알파벳 하나하나씩 찾아가며, 마지막에 * 가 있는지 확인한다.