[Programmers Lv.4] 자동완성 - Python

문제

포털 다음에서 검색어 자동완성 기능을 넣고 싶은 라이언은 한 번 입력된 문자열을 학습해서 다음 입력 때 활용하고 싶어 졌다. 예를 들어, go 가 한 번 입력되었다면, 다음 사용자는 g 만 입력해도 go를 추천해주므로 o를 입력할 필요가 없어진다! 단, 학습에 사용된 단어들 중 앞부분이 같은 경우에는 어쩔 수 없이 다른 문자가 나올 때까지 입력을 해야 한다.효과가 얼마나 좋을지 알고 싶은 라이언은 학습된 단어들을 찾을 때 몇 글자를 입력해야 하는지 궁금해졌다.
예를 들어, 학습된 단어들이 아래와 같을 때
go gone guild
- go를 찾을 때 go를 모두 입력해야 한다.
- gone을 찾을 때 gon 까지 입력해야 한다. (gon이 입력되기 전까지는 go 인지 gone인지 확신할 수 없다.)
- guild를 찾을 때는 gu 까지만 입력하면 guild가 완성된다.
이 경우 총 입력해야 할 문자의 수는 7이다.
라이언을 도와 위와 같이 문자열이 입력으로 주어지면 학습을 시킨 후, 학습된 단어들을 순서대로 찾을 때 몇 개의 문자를 입력하면 되는지 계산하는 프로그램을 만들어보자.

 

효율적인 검색 알고리즘을 위해 트라이 자료구조를 이용해 풀었다.

 

트라이(Trie)는 문자열을 저장하고 효율적으로 탐색하기 위한 트리 형태의 자료구조로, 자동완성 기능이나 사전 검색 등 문자열을 탐색하는데 특화되어있다.

이진 검색 트리의 경우 일반적으로 O(logn)의 시간 복잡도를 가지지만, 문자열의 경우 두 문자열을 비교하기 위해서는 문자열의 길이만큼 시간이 걸리기 때문에 원하는 문자열을 찾기 위해서는 O(mlogn)의 시간이 걸린다. 이 문제를 해결하기 위한 것이 트라이다.

 

문제 해결을 위해 Node와 Trie class를 만들었다. 트리의 각 노드가 {key, data, children}으로 구성되도록 했는데, key는 문자, data는 해당 노드에서 완성되는 문자열, children은 다음 문자를 담은 노드를 저장한다.

 

다만 이 문제에서는 추가적인 조건이 필요했다. 보통 트라이에서 문자열을 탐색할 때 문자열의 마지막 문자까지 탐색한 후 해당 노드의 data에 찾는 문자열이 있으면 탐색에 성공했다고 보는데, 문자열을 끝까지 찾는게 아니라 문자열을 찾을 수 있는 최소 글자 수를 구하는 문제이기 때문이다.

어떻게 할 지 고민하다가 각 노드에 앞으로 남은 단어 개수가 몇개인지, 즉 모든 children 중 data 필드가 있는 Node가 몇 개인지를 저장하기로 했다. 어떤 노드에 노달했을 때 앞으로 완성될 단어가 하나 뿐이라면 탐색을 그만해도 되기 때문이다.

남은 단어 개수는 Trie에 단어를 추가할 때 지나는 노드에 1씩 해당 값을 더해주면서 카운트했다.

 

코드

class Node:
    def __init__(self, key, data=None):
        self.key = key
        self.data = data
        self.children = {}
        self.remain = 0

class Trie:
    def __init__(self):
        self.root = Node(None)
        
    def add(self, str):
        cur_node = self.root
        for c in str:
            if c not in cur_node.children:
                cur_node.children[c] = Node(c)
            cur_node = cur_node.children[c]
            cur_node.remain += 1
        cur_node.data = str
    
    def search(self, str):
        cur_node = self.root
        res = 0
        for c in str:
            if cur_node.remain == 1:
                break
            if c in cur_node.children:
                res += 1
                cur_node = cur_node.children[c]
        return res

def solution(words):
    answer = 0
    trie = Trie()
    for w in words:
        trie.add(w)
    for w in words:
        answer += trie.search(w)
    return answer