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

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

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

 

저번에 풀었던 [3차] 자동완성에서 썼던 Trie 자료구조를 이용하였다.

 

'?' 가 접두사 아니면 접미사로만 존재하기 때문에 두 개의 dictionary 를 만들어 검색할 수 있도록 하였다.

예를 들어, 'frodo' 를 dictionary에 추가시키면 

{'f': {'*': [5], 'r': {'*': [5], 'o': {'*': [5], 'd': {'*': [5], 'o': {'*': [5]}}}}}} 이고 '*' 는 그 위치의 문자를 가진 문자열의 길이가 5인 것이 하나 존재한다는 의미이다.

 

'front', 'frost' 과 'fronzen'를 추가시키면, 

{'f': {'*': [5, 5, 5, 6], 'r': {'*': [5, 5, 5, 6], 'o': {'*': [5, 5, 5, 6], 'd': {'*': [5], 'o': {'*': [5]}}, 'n': {'*': [5], 't': {'*': [5]}}, 's': {'*': [5], 't': {'*': [5]}}, 'z': {'*': [6], 'e': {'*': [6], 'n': {'*': [6]}}}}}}} 로

'f' 로 시작하는 문자열은 길이가 5인 것이 세 개, 6인 것이 한 개 존재한다는 것이고, 'f' -> 'r' -> 'o' -> 'z' 에서 '*' 의 값이 [6] 이기 때문에 접두사가 'froz' 로 해당하는 문자열은 길이가 6인 것이 한 개 존재한다는 것이다.

 

즉, 완성한 dictionary 를 통해 'fro??' 와 일치하는 문자열을 찾을 때는 '?' 가 나오기 전에 'f' -> 'r' -> 'o' 에서의 '*' 에 해당하는 배열의 원소 중에서 'fro??' 의 길이인 5의 개수를 구하면 된다. 

 

class Trie:
    dictionary = {}
    dictionaryReverse = {}
    length = []
    
    def insert(self, string):
        current_dic = self.dictionary
        current_dicReverse = self.dictionaryReverse
        N = len(string)        
        self.length.append(N)
        
        for i in range(N):
            C = string[i]
            R = string[N-1-i]
            
            if C not in current_dic:
                current_dic[C] = {'*' : [N]}
            else:
                current_dic[C]['*'].append(N)
            
            current_dic = current_dic[C]
            
            if R not in current_dicReverse:
                current_dicReverse[R] = {'*' : [N]}
            else:
                current_dicReverse[R]['*'].append(N)
            
            current_dicReverse = current_dicReverse[R]

    def search(self, string):
        current_dic = self.dictionary if string[0] != '?' else self.dictionaryReverse
        string = string if string[0] != '?' else string[::-1]
        N = len(string)
        
        if string[0] == '?':
            return self.length.count(N)
        else:
            for i in string:
                if i in current_dic:
                    current_dic = current_dic[i]
                else:
                    if i != '?':
                        return 0
                    else:
                        return current_dic['*'].count(N)
        
def solution(words, queries):
    answer = []
    T = Trie()
    
    for i in words:
        T.insert(i)
    
    for i in queries:
        answer.append(T.search(i))

    return answer

print(solution(['frodo', 'front', 'frost', 'frozen', 'frame', 'kakao'], ['fro??', '????o', 'fr???', 'fro???', 'pro?'])) # [3, 2, 4, 1, 0]