[프로그래머스] 가사 검색 (Python)

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

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

문자열 검색 알고리즘인 Trie를 활용해서 푸는 문제다

계속 효율성 4번에서 시간 초과가 나서 풀이를 참고하고 그대로 구현했다.

 

class Trie():
    def __init__(self):
        self.child = dict()
        self.count = 0
    
    def insert(self,str):
        curr = self
        for ch in str:
            curr.count += 1
            if ch not in curr.child:
                curr.child[ch] = Trie()
            curr = curr.child[ch]
        curr.count += 1
        
    def search(self,str):
        curr = self
        for ch in str:
            if ch == '?':
                return curr.count
            if ch not in curr.child:
                return 0
            else:
                curr = curr.child[ch]
        return curr.count
    
def solution(words, queries):
    answer = []
    TrieRoot = [Trie() for _ in range(10000)]
    ReTrieRoot = [Trie() for _ in range(10000)]
    
    for word in words:
        TrieRoot[len(word)-1].insert(word)
        ReTrieRoot[len(word)-1].insert(word[::-1])
    
    for query in queries:
        if query[0] != '?':
            answer.append(TrieRoot[len(query)-1].search(query))
        else:
            answer.append(ReTrieRoot[len(query)-1].search(query[::-1]))
    return answer

이번 기회에 class를 활용하는 문제에 좀 자신감이 생긴 것 같다