[프로그래머스] 자동완성 - python

자동완성

문제 설명

포털 다음에서 검색어 자동완성 기능을 넣고 싶은 라이언은 한 번 입력된 문자열을 학습해서 다음 입력 때 활용하고 싶어 졌다. 예를 들어, go 가 한 번 입력되었다면, 다음 사용자는 g 만 입력해도 go를 추천해주므로 o를 입력할 필요가 없어진다! 단, 학습에 사용된 단어들 중 앞부분이 같은 경우에는 어쩔 수 없이 다른 문자가 나올 때까지 입력을 해야 한다. 효과가 얼마나 좋을지 알고 싶은 라이언은 학습된 단어들을 찾을 때 몇 글자를 입력해야 하는지 궁금해졌다.

예를 들어, 학습된 단어들이 아래와 같을 때

go
gone
guild

• go를 찾을 때 go를 모두 입력해야 한다.
• gone을 찾을 때 gon 까지 입력해야 한다. (gon이 입력되기 전까지는 go 인지 gone인지 확신할 수 없다.)
• guild를 찾을 때는 gu 까지만 입력하면 guild가 완성된다.

이 경우 총 입력해야 할 문자의 수는 7이다.

라이언을 도와 위와 같이 문자열이 입력으로 주어지면 학습을 시킨 후, 학습된 단어들을 순서대로 찾을 때 몇 개의 문자를 입력하면 되는지 계산하는 프로그램을 만들어보자.

입력 형식

학습과 검색에 사용될 중복 없는 단어 N개가 주어진다. 모든 단어는 알파벳 소문자로 구성되며 단어의 수 N과 단어들의 길이의 총합 L의 범위는 다음과 같다.

• 2 <= N <= 100,000
• 2 <= L <= 1,000,000

출력 형식

단어를 찾을 때 입력해야 할 총 문자수를 리턴한다.

풀이

문자열을 빠르게 탐색하기 위해서 Trie 자료구조를 이용할 수 있다. Trie 자료구조의 각 노드는 이전 노드로부터 온 문자열을 더한 값을 가지고 있다.

 

Trie 자료구조를 만들어 단어를 입력하고, 루트부터 각 단어를 탐색하면서 Node의 갯수가 1이 될 때를 리턴하면 된다. (Node의 개수가 1이라는 것은 중복되는 단어가 없다는 뜻이므로 )

class Node:
    def __init__(self):
        self.children = {}
        self.cnt = 0
​
class Trie:
    def __init__(self):
        self.root = Node()
    
    def insert(self, string):
        cur_node = self.root
        for char in string:
            cur_node.children.setdefault(char, Node())
            cur_node.children[char].cnt += 1
            cur_node = cur_node.children[char]
            
    def auto_comp(self, string):
        cur_node = self.root
        for i, char in enumerate(string, 1):
            if cur_node.children[char].cnt == 1:
                return i
            cur_node = cur_node.children[char]
        return len(string)
        
def solution(words):
    answer = 0
    trie = Trie()
    for word in words:
        trie.insert(word)
    
    for word in words:
        answer += trie.auto_comp(word)
    return answer

이렇게 자료구조를 다 만들지 않고 딕셔너리를 사용해 간단하게 풀 수도 있다.

def make_trie(words):
    root = {}
    for word in words:
        cur_node = root
        for char in word:
            cur_node.setdefault(char, [0,{}])
            cur_node[char][0] += 1
            cur_node = cur_node[char][1]
    return root
​
def solution(words):
    answer = 0
    trie = make_trie(words)
    
    for word in words:
        cur_node = trie
        for i, char in enumerate(word, 1):
            if cur_node[char][0] == 1:
                break
            cur_node = cur_node[char][1]
        answer += i # 다 쳐야되는 경우 , 부분만 쳐야되는 경우 포함 (go, gone)
    
    return answer