가사 검색

from collections import defaultdict, Counter import sys sys.setrecursionlimit(10**9) class Node(object): def __init__(self, key, data=None): self.key = key self.data = data self.n = 0 self.children = {} def finalize(self): self.n = 1 if self.data is not None else 0 for char, node in self.children.items(): self.n += node.finalize() return self.n class Trie(object): def __init__(self): self.head = Node(None) def insert(self, string): curr_node = self.head for char in string: if char not in curr_node.children: curr_node.children[char] = Node(char) curr_node = curr_node.children[char] curr_node.data = string def finalize(self): self.head.finalize() def search(self, string): curr_node = self.head for char in string: if char == '?': break if char in curr_node.children: curr_node = curr_node.children[char] else: return 0 return curr_node.n def solution(words, queries): prefix, suffix = defaultdict(Trie), defaultdict(Trie) for word in words: l = len(word) prefix[l].insert(word) suffix[l].insert(word[::-1]) for l, t in prefix.items(): t.finalize() for l, t in suffix.items(): t.finalize() lengthcount = Counter([len(w) for w in words]) answer = [] for q in queries: l = len(q) # 예외처리) q가 모두 ?인 경우는 같은 길이를 갖는 단어의 개수를 리턴합니다. if q[0] == '?' and q[-1] == '?': answer.append(lengthcount[l]) continue if q[0] == '?': # suffix answer.append(suffix[l].search(q[::-1])) else: answer.append(prefix[l].search(q)) return answer