가사 검색

• 트라이 (Trie)를 사용하면 간단히 풀리는 문제.
• 생각할 부분은 단순히 특정 접두어로 시작하는 단어의 갯수가 아니라, 특정 접두어로 시작하면서 전체 길이가 x인 단어의 갯수를 구해야 한다는 것.
  • 트라이 구조를 조금 수정해서, 트라이의 각 노드에 차일드의 갯수를 길이별로 구분해서 저장한는 방법이 직관적이긴 하다.
  • 그러나 트라이 구조를 수정 않고, 라이브러리의 트라이를 그대로 활용하고 싶으면, 무식해보이지만 단어 길이별로 각기 다른 트라이를 n개 만들어서 (길이 5짜리 단어를 담는 트라이, 길이 6짜리 단어를 담는 트라이, …, 이런식으로) 처리하는 방법도 가능하다. 실제로도 이 방법으로 해결함.
• ?가 뒤가 아니라 앞에 붙어서, 특정 접미사로 끝나는 단어의 갯수를 세야 할 경우도 마찬가지로, string을 reverse한 것들을 저장하는 트라이를 만들면 된다.
• 트라이를 생성하는 데에는 O(w), 쿼리를 처리하는 데에는 O(q). 그래서 총 시간복잡도는 O(w+q). (w=전체 가사 단어 길이의 합, q=전체 쿼리 길이의 합)

"""Solution code for "Programmers 60060. 가사 검색".

- Problem link: https://programmers.co.kr/learn/courses/30/lessons/60060
- Solution link: http://www.teferi.net/ps/problems/programmers/60060
"""

from teflib import string

MAX_LEN = 10000


def solution(words, queries):
    tries = [string.Trie() for _ in range(MAX_LEN + 1)]
    rev_tries = [string.Trie() for _ in range(MAX_LEN + 1)]
    for word in words:
        tries[len(word)].add(word)
        rev_tries[len(word)].add(reversed(word))

    answer = []
    for query in queries:
        try:
            if query[0] == '?':
                node = rev_tries[len(query)][reversed(query.lstrip('?'))]
                answer.append(rev_tries[len(query)].get_child_count(node))
            else:
                node = tries[len(query)][query.rstrip('?')]
                answer.append(tries[len(query)].get_child_count(node))
        except KeyError:
            answer.append(0)

    return answer

• Dependency: teflib.string.Trie