[2020 KAKAO BLIND RECRUITMENT] 가사 검색

내가 푼 풀이

class Node(object):
    def __init__(self, key):
        self.key = key
        self.count = 0
        self.children = {}

class Trie(object):
    def __init__(self):
        self.head = Node(self)
    
    def insert(self, string):
        curr_node = self.head

        for char in string:
            curr_node.count += 1
            if char not in curr_node.children:
                curr_node.children[char] = Node(char)
            curr_node = curr_node.children[char]
        curr_node.count += 1

    def starts_with(self, prefix):
        curr_node = self.head

        for char in prefix:
            if char == '?':
                break
            
            if char in curr_node.children:
                curr_node = curr_node.children[char]
            else :
                return 0
        return curr_node.count

def solution(words, queries):
    answer = []
    tries = {}
    reverse_tries = {}

    for word in words:
        if len(word) in tries:
            tries[len(word)].insert(word)
            reverse_tries[len(word)].insert(reversed(word))
        else:
            trie = Trie()
            reverse_trie = Trie()

            trie.insert(word)
            reverse_trie.insert(reversed(word))

            tries[len(word)] = trie
            reverse_tries[len(word)] = reverse_trie

    for query in queries:
        if len(query) in tries:
            if query[0] != '?':
                trie = tries[len(query)]
                answer.append(trie.starts_with(query))
            else:
                trie = reverse_tries[len(query)]
                answer.append(trie.starts_with(reversed(query)))
        else:
            answer.append(0)
            
    return answer

이 문자는 문자열 처리 알고리즘인 Trie를 사용하는 문제였다. 이 알고리즘을 사용하지 않으면 효율성에서 통과를 못한다고 들었다!
node를 이용하여 트리구조로 문자열을 나열하는 방법인데, word가 중복을 허용하지 않으므로 트리를 만들때부터 count를 세어 넣어주었다.

word의 길이에 따라 개수가 달라지므로 trie가 길이별로(len(word)) 필요하며, '?'가 맨 앞에 올수도 있으므로 역순으로도 trie를 만들어야 한다. trie를 알고있다면 어렵지는 않은 문제였지만, 나는 몰랐었기 때문에(..^^) 더 열심히 공부해야겠따~~