가사 검색

문제 링크 : programmers.co.kr/learn/courses/30/lessons/60060?language=python3

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

가사(words)의 최대 개수가 100,000개, 검색 키워드(queries)의 최대 개수가 100,000개이므로 

모든 검색 키워드에 대해서 가사와 비교하면 많은 시간이 소요됩니다.

 

 

효율적인 검색을 위해 Trie 자료 구조를 이용했습니다.

Trie는 여러 문자열을 저장하고 있고 입력으로 임의의 문자열이 주어졌을 때 Trie에 포함된 것인지 빠르게 확인할 수 있습니다. 찾고자 하는 문자열의 길이가 m이라면 Trie는 O(m) 만큼의 시간복잡도를 갖습니다.

 

 

입력으로 주어지는 가사(words)의 각 문자를 '?'로 치환하여 Trie에 저장할 수 있습니다.

예를 들어 words에 "one"가 입력으로 주어졌다면

Trie에는 "?ne", "??e", "???', "on?", "o??"의 5가지 경우를 저장합니다.

(문제에서 검색 키워드는 하나 이상의 '?'를 포함하고 '?'가 중간에 있거나 양쪽으로 있는 경우는 존재하지 않는다고 명시하고 있습니다.)

 

 

저는 물음표의 위치가 1) 앞에서부터 시작하는 경우, 2) 뒤에서부터 시작하는 경우로 나누어 각각 dictionary에 저장했습니다. (각각 ahead, behind로 정의했습니다.)

위의 예시에서

ahead에는 "?ne", "??e"

behind에는 "???", "on?", "o??"

가 저장됩니다.

 

 

검색을 할 때는 물음표가 앞에서부터 시작하는지 판별한 뒤에

물음표가 앞에서부터 시작한다면 ahead에서

모든 문자열이 물음표로 만들어졌거나 물음표가 뒤에서부터 시작한다면 behind에서 탐색을 진행했습니다.

 

 

class Trie:
    ahead = dict()
    behind = dict()

    def add(self, s:str):
        cur = self.behind
        for i in range(len(s)):
            ch = s[i]
            LEN = len(s) - i
            if not cur.__contains__(ch):
                cur.update({ch:dict()})
            if not cur.__contains__(LEN):
                cur.update({LEN:0})
            cur[LEN] += 1
            cur = cur[ch]
        
        cur = self.ahead
        for i in range(len(s)-1, -1, -1):
            ch = s[i]
            LEN = i + 1
            if not cur.__contains__(ch):
                cur.update({ch:dict()})
            if not cur.__contains__(LEN):
                cur.update({LEN:0})
            cur[LEN] += 1
            cur = cur[ch]

    def search(self, query:str) -> int:
        if query[0] == '?':
            cur = self.ahead
            idx = query.rfind('?')
            cnt = idx + 1
            for i in range(len(query)-1, -1, -1):
                ch = query[i]
                if ch == '?':
                    break
                elif not cur.__contains__(ch):
                    return 0
                cur = cur[ch]
            return cur[cnt] if cur.__contains__(cnt) else 0
        else:
            cur = self.behind
            idx = query.find('?')
            cnt = len(query) - idx
            for i in range(idx):
                ch = query[i]
                if not cur.__contains__(ch):
                    return 0
                cur = cur[ch]
            return cur[cnt] if cur.__contains__(cnt) else 0            

def solution(words, queries):
    answer = []
    trie = Trie()
    
    for word in words:
        trie.add(word)

    for query in queries:
        answer.append(trie.search(query))

    return answer