2018 KAKAO 자동완성

https://programmers.co.kr/learn/courses/30/lessons/17685

코딩테스트 연습 - [3차] 자동완성

 

이 문제를 풀기 전에 유사한 문제인 KAKAO 2020 가사검색 문제를 풀었기 때문에 금방 풀 수 있었다.
보자마자 Trie로 접근해야겠다는 생각이 들었고, 통과를 할 수 있었다.

Trie와 다른 점은 현재 node 를 root로 하는 subtree에 몇개의 단어가 저장되어있는지 알 수 있는 cnt 변수를 추가한 것이다.

원래 각 단어당 입력해야 하는 문자의 개수를 child node를 저장하는 next의 길이로 판단하려고 했다.
하지만 [go, gone] 과 같은 입력에서 등록되어있는 child node가 g 밖에 없어 1만 반환 되기 때문에 정답을 반환할 수 없어 cnt 변수를 추가했다.

space (' ')는 문자의 마지막 Null 대신 넣었다. 이를 추가함으로써 child node에서 단어의 개수를 셀 때 현재 node에서 종료되는 단어의 처리를 따로 해주지 않게 되었다.

import sys
sys.setrecursionlimit(1000100)

class trie:
    def __init__(self):
        self.next = dict()
        self.cnt = 0
    
    def insert(self, s, i):
        self.cnt += 1
        c = s[i]
        if c not in self.next:
            self.next[c] = trie()
        if c != ' ':
            self.next[c].insert(s, i+1)
        
    def search(self, s, i):
        c = s[i]
        if c == ' ': return i
        if self.cnt < 2: return i
        return self.next[c].search(s, i+1)
        

def solution(words):
    answer = 0
    t = trie()
    for w in words:
        t.insert(w+' ', 0)
    
    for w in words:
        answer += t.search(w+' ', 0)
        pass
    
    return answer