[programmers] Level4. 가사검색

풀이

• Trie 자료구조 사용하지 않으면 효율성 테스트에서 시간초과.
  • 이전에 map, filter로 풀어봤지만 효율성 테스트 1,2에서 시간초과발생했음.
• wildcard가 처음에 나오는 경우, 끝에서 나오는경우를 생각해야 함.
  • words를 원래 순서대로 넣은 Trie, 거꾸로 넣은 Trie 2개를 활용.
• 문제 조건에 중복된 word나 query가 있을 수 있기 때문에, 각각 dictionary를 사용해 이전 값을 저장.

class Node(object):
    def __init__(self, w=None):
        # 알파벳 하나를 저장
        self.word = w
        # 다음 등장하는 모든 알파벳을 저장
        self.children = {}
        # 단어 끝 알파벳에 단어 길이를 저장(search연산에서 dfs의 종료조건으로 활용)
        self.len = 0

class Trie:
    def __init__(self):
        self.root = Node(None)
        # 중복 가사를 처리하기 위한 dictionary
        # key : word, value : len(word)
        self.keyList = {}

    def insert(self, word):
        curr = self.root
        # 중복되지 않았을 경우에만 삽입 연산 수행함. 
        if word not in self.keyList.keys():
            self.keyList[word] = len(word)
            for w in word:
                # 알파벳이 자식 노드에 없으면 추가 후 이동
                if w not in curr.children.keys():
                    curr.children[w] = Node(w)
                curr = curr.children[w]
            # 마지막에 가사 길이를 저장. 나머지 알파벳 노드는 초기화 값 0이 저장되어있음.
            curr.len = len(word)

    def search(self, word):
        num = 0
        curr = self.root
        num_wildcard = word.count('?')
        word_len = len(word)
        # wildcard 앞에까지 curr를 이동.
        for w in word[:word_len - num_wildcard]:
            if w in curr.children.keys():
                curr = curr.children[w]
            # 하지만 그 전에 알파벳이 없으면 0을 리턴
            else:
                return num

        stack = []
        for node in curr.children.values():
            # curr, depth 저장한 stack으로 dfs연산 수행
            stack.append((node, 1))

        while stack:
            curr, depth = stack.pop()
            if curr.len == word_len and depth == num_wildcard:
                num += 1

            if depth < num_wildcard:
                for node in curr.children.values():
                    stack.append((node, depth+1))

        return num

def solution(words, queries):
    answer = []
    dic = {}
    trie = Trie()
    reverse_trie = Trie()
    for w in words:
        trie.insert(w)
        reverse_trie.insert(w[::-1])
    for q in queries:
        if q not in dic:
            if q[0] == '?':
                answer.append(reverse_trie.search(q[::-1]))
            else:
                answer.append(trie.search(q))
            dic[q] = answer[-1]
        else:
            answer.append(dic[q])
    return answer