title: "[프로그래머스] 가사 검색 Python 파이썬 해설 (Level 4) - 이도훈"
cleanUrl: "programmers/60060"
description: "프로그래머스 Level 4 문제 [가사 검색]의 풀이를 정리합니다."

문제 설명 및 제한사항

아이디어 및 해결 방법

• Prefix/Suffix 별, 단어 길이 별로 Trie를 구성하여 해결합니다.

코드

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

출처

프로그래머스 코딩테스트 연습 https://school.programmers.co.kr/learn/challenges