[프로그래머스] 가사 검색 / 2020 KAKAO BLIND RECRUITMENT / python

트라이 자료구조를 사용하여 문제를 풀면 효율성 테스트를 통과할 수 있다.
트라이(Trie) 자료구조

와일드카드 문자에 해당하는 가사 단어가 트라이 자료구조에 몇개가 존재하는지 알기 위해서 트라이 자료구조에 가사 단어를 insert 할 때 각 알파벳마다 존재하는 count에 1씩 더해준다. 이 때 와일드카드 문자가 모두 '?'일 수도 있으므로 트라이 자료구조 head의 count도 계속 1씩 더해준다.

이제 트라이 자료구조에서 와일드카드 문자에서 알파벳에 해당하는 부분의 마지막 노드를 찾아서 count를 반환하면 될까 싶지만 그렇지 않다.
'fro???'와 'fro?????' 모두 트라이 자료구조에서 알파벳 o 노드에 해당하는 count를 반환하기 때문이다.

그렇기 때문에 단어의 길이마다 각각 트라이 자료구조를 만들어줘야 한다.

조건에 가사 길이는 1 이상 10000 이하라고 나와 있으므로 10000개의 트라이 자료구조를 만들어주면 된다.
그리고 문제 조건에서 '?'는 접두사나 접미사로 주어진다고 했으므로 와일드카드 문자를 뒤집은 다음에 트라이 자료구조를 사용해야함으로 거꾸로된 가사 단어들의 트라이 자료구조도 필요하다.
결과적으로 총 20000개의 트라이 자료구조가 필요하다.

와일드카드 문자를 확인할 때는 '?'가 접두사나 접미사 중 어디로 주어졌는지(트라이 자료구조와 뒤집어진 트라이 자료구조 중 어느걸 사용할지)와 와일드카드 문자의 길이를 사용하여 알맞은 트라이 자료구조를 고르고 해당 자료구조에서 와일드카드 문자에서 알파벳에 해당하는 부분의 마지막 노드를 찾아서 count를 반환하면 된다.

class Node:
    def __init__(self, data):
        self.data = data
        self.count = 0
        self.child = {}


class Trie:
    def __init__(self):
        self.head = Node(None)

    def insert(self, string):
        cur = self.head
        cur.count += 1

        for c in string:
            if c not in cur.child:
                cur.child[c] = Node(c)
            cur = cur.child[c]
            cur.count += 1

    def count(self, prefix):
        cur = self.head

        for c in prefix:
            if c not in cur.child:
                return 0
            cur = cur.child[c]

        return cur.count


def solution(words, queries):
    answer = []

    tries = create_trie(words)
    reversed_tries = create_trie(words, True)

    for query in queries:
        answer.append(count_matched_word(tries, reversed_tries, query))

    return answer


def create_trie(words, is_reversed=False):
    trie_dic = {i: Trie() for i in range(1, 10001)}

    for word in words:
        if is_reversed:
            word = word[::-1]
        trie_dic[len(word)].insert(word)

    return trie_dic


def count_matched_word(tries, reversed_tries, query):
    no_mark_query = query.replace('?', '')

    if query[0] == '?':
        return reversed_tries[len(query)].count(no_mark_query[::-1])
    else:
        return tries[len(query)].count(no_mark_query)