카카오 가사검색

from collections import deque, defaultdict
class Node(object):
    def __init__(self, key, data=None):
        self.key = key
        self.data = data
        self.count = 0
        self.child = {}
class Trie():
    def __init__(self):
        self.head = Node(None)

    def insert(self, string):
        cursor = self.head
        for s in string:
            if s not in cursor.child:
                cursor.child[s] = Node(s)
            cursor.child[s].count += 1
            cursor = cursor.child[s]
        cursor.data = string

    def search(self, string):
        cursor = self.head
        for s in string:
            if s in cursor.child:
                cursor = cursor.child[s]
            else:
                return False

        if cursor.data == string:
            return True
        else:
            return False

    def starts_with(self, prefix, num):
        cursor = self.head
        ret = []
        for s in prefix:
            if s in cursor.child:
                cursor = cursor.child[s]
            else:
                return 0
        return cursor.count

def solution(words, queries):
    answer = []
    myTrie = {}
    reversedTrie = {}
    wl = defaultdict(int)
    for word in words:
        wordLen = len(word)
        wl[wordLen] += 1
        if wordLen not in myTrie:
            myTrie[wordLen] = Trie()
            reversedTrie[wordLen] = Trie()
        myTrie[wordLen].insert(word)
        reversedTrie[wordLen].insert(word[::-1])

    for query in queries:
        queryLen = len(query)
        cnt=0
        if query.startswith('?'):
            query = query[::-1]
            for i in range(len(query)):
                if query[i] == '?':
                    break
                cnt+=1
            if cnt == 0 :
                answer.append(wl[queryLen])
            else:
                if queryLen in reversedTrie:
                    answer.append(reversedTrie[queryLen].starts_with(query[:cnt], queryLen))
                else :
                    answer.append(0)
        else:
            for i in range(len(query)):
                if query[i] == '?':
                    break
                cnt += 1
            if queryLen in myTrie:
                answer.append(myTrie[queryLen].starts_with(query[:cnt], queryLen))
            else :
                answer.append(0)

    return answer

print(solution(["frodo", "front", "frost", "frozen", "frame", "kakao"], ["fro??", "????o", "fr???", "fro???", "pro?"]))

평범한 Trie로는 못풀고 count 값을 줘야한다. 그리고 단어 갯수별로 Trie를 나누지 않으면 count가 제대로 먹히지 않는다.

 

ex) 나누지 않는다면 starts_with 에서 fro 로 시작하는 모든 단어의 갯수가 리턴됨. 이를 방지하기 위해 len == 5 와 같은 로직이 들어가야하는데 이러면 무조건 시간초과