[프로그래머스] 가사 검색 / Python / 2020 KAKAO BLIND RECRUITMENT

programmers.co.kr/learn/courses/30/lessons/60060?language=python3

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

---

---

Solution

class Node:
    def __init__(self, data=None):
        self.data = data
        self.child = {}
        self.length = {}

class Trie:
    def __init__(self):
        self.head = Node(None)

    def insert(self, string):
        curr_node = self.head
        length = len(string)

        if not length in curr_node.length:
            curr_node.length[length] = 1
        else:
            curr_node.length[length] += 1

        for char in string:
            if not char in curr_node.child:
                curr_node.child[char] = Node(char)
            curr_node = curr_node.child[char]
            length -= 1
            if not length in curr_node.length:
                curr_node.length[length] = 1
            else:
                curr_node.length[length] += 1

    def search(self, query, wild):
        curr_node = self.head

        for char in query:
            if char in curr_node.child:
                curr_node = curr_node.child[char]
            else:
                return 0

        if not wild in curr_node.length:
            return 0
        else:
            return curr_node.length[wild]


def solution(words, queries):
    pre_trie = Trie()
    inv_trie = Trie()

    for i in words:
        pre_trie.insert(i)
        inv_trie.insert(i[::-1])

    result = []
    for q in queries:
        if q[0] != '?':
            question = q.find('?')
            query = q[:question]
            wild = len(q[question:])
            result.append(pre_trie.search(query, wild))
        else:
            q = q[::-1]
            question = q.find('?')
            query = q[:question]
            wild = len(q[question:])
            result.append(inv_trie.search(query, wild))

    return result

 

트라이 자료구조를 이용해 푸는 문제이다. 트라이는 다른 분의 포스팅을 참고했다.

처음에는 Node의 length를 배열로 선언하고, 남은 스트링의 개수를 배열에 넣고 확인하는 방식으로 구현을 했었다. 하지만 효율성 2번과 3번을 통과하지 못했고 length를 dict로 선언하는 방식으로 바꿔서 풀었다. 먼저, Trie의 insert메서드에서 string의 각 글자를 돌며 각 Node에 남은 length값을 넣어준다. 같은 값의 length가 이미 있으면 개수를 1더해준다.

 

    def insert(self, string):
        curr_node = self.head
        length = len(string)

        if not length in curr_node.length:
            curr_node.length[length] = 1
        else:
            curr_node.length[length] += 1

        for char in string:
            if not char in curr_node.child:
                curr_node.child[char] = Node(char)
            curr_node = curr_node.child[char]
            length -= 1
            if not length in curr_node.length:
                curr_node.length[length] = 1
            else:
                curr_node.length[length] += 1

 

물음표가 앞에서부터 있을 수 있고, 뒤에서부터 있을 수 있기 때문에 Trie를 순방향과 역방향 두 개 만들어준다. 그리고 queries를 확인하며 물음표가 뒤에 있는 경우 물음표의 인덱스를 찾아 물음표의 개수(wild)와 물음표를 뺀 문자열(query)를 선언한다. 물음표부터 앞에 있는 경우의 query는 query를 뒤집고 똑같은 로직으로 wild와 query를 선언한다. 그 후 search메서드에 인자로 넘겨준다. 

 

    for q in queries:
        if q[0] != '?':
            question = q.find('?')
            query = q[:question]
            wild = len(q[question:])
            result.append(pre_trie.search(query, wild))
        else:
            q = q[::-1]
            question = q.find('?')
            query = q[:question]
            wild = len(q[question:])
            result.append(inv_trie.search(query, wild))