가사 검색

programmers.co.kr/learn/courses/30/lessons/60060

코딩테스트 연습 - 가사 검색

from collections import defaultdict

class Node:
    def __init__(self, w=None):
        self.val = w # 노드 값
        self.count = 0 # 현재 노드에서 이어지는 단어수
        self.children = defaultdict(Node) # 자식 노드
class Trie:
    def __init__(self):
        self.head1 = Node() # 단어 정방향 저장 트라이
        self.head2 = Node() # 단어 역방향 저장 트라이
        
    def insert(self, word): # 트라이에 단어 삽입
        cur, cur_rvs, l = self.head1, self.head2, len(word) 
        cur, cur_rvs = cur.children[l], cur_rvs.children[l] # 트라이는 단어 글자 수로 우선 분류
        cur.count += 1 
        cur_rvs.count += 1
        for i in range(l): # 한 문자씩 트라이 구성
            if word[i] not in cur.children:
                cur.children[word[i]] = Node(word[i])
            if word[l - 1 - i] not in cur_rvs.children:
                cur_rvs.children[word[l - 1 - i]] = Node(word[l - 1 - i])
            cur, cur_rvs = cur.children[word[i]], cur_rvs.children[word[l - 1 - i]]
            cur.count += 1 # 각 문자 노드에서 앞으로 이어질 단어 개수 저장
            cur_rvs.count += 1
            
    def search(self, word): # 트라이 단어 검색
        cur = self.head1 if word[-1] == '?' else self.head2 # 검색 단어의 '?'위치에 따라 트라이 선택
        word = word if word[-1] == '?' else word[::-1] # 역방향 트라이 이용시 단어 뒤집기
        l = len(word)
        if cur.children[l].count == 0: return 0 # 해당 단어 크기가 존재 하지 않으면 바로 0 반환
        cur = cur.children[l] # 해당 단어 크기 노드로 이동
        for i in range(l): # 조건에 맞게 노드 탐색
            if word[i] == '?': # '?'가 나오면 현재 노드의 count를 반환
                return cur.count
            if word[i] not in cur.children: # 일치하는 문자가 없으므로 0 반환
                return 0
            cur = cur.children[word[i]] # 다음 문자 노드로 이동
                
def solution(words, queries):
    answer = []
    T = Trie()
    
    for word in words: # 트라이 구성
        T.insert(word)
    
    for query in queries: # 단어 검색
        answer.append(T.search(query))
    
    return answer

→ 이전에 자동 완성 문제와 비슷하여 같은 자료구조 트라이를 사용함.

처음에는 head에서 단어 크기로 나누고 이후에 트라이를 구성하는 방식이 아니라,

각 노드 마다 {남은 단어 크기 : 개수}를 저장하는 딕셔너리를 사용해 보았는데 

위와 시간적으로 큰 차이는 없을것 같지만 효율성에서 두개에 걸렸다.