프로그래머스 - 가사검색- Trie자료구조

09 Sep 2020

문제

입력

---

나의 처음 코드

처음에는 startswith, endswith의 기능을 사용하였지만 효율성 2번이 통과하지 못해 실패했다. 또한 딕셔너리를 이용하여 푸는 풀이도 효율성이 실패했다. 다음으로 Trie 자료구조를 사용하였다. 이진트리와는 다르게 자식들이 여럿 있을 수 있으므로, 부모가 자식을 저장하는 자료구조를 잘 설계하는 것이 중요하다. 또한, 리프노드에서 일치하는 단어를 찾으면 시간초과가 날 수 있으므로, 각 노드별로 가지고 있는 자식들의 길이 종류? (자식들을 길이별로 갯수를 체크한) 자료구조를 사용하면, 리프노드까지 안 가도 빠르게 셀 수 있다.

---

풀이 사진

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

def insert(word, head):
    curNode = head
    w_len = len(word)
    for w in word:
        if w in curNode.children:
            curNode = curNode.children[w]
        else:
            newNode = Node(w)
            curNode.children[w] = newNode
            curNode = curNode.children[w]
        if w_len in curNode.length:
            curNode.length[w_len] += 1
        else:
            curNode.length[w_len] = 1
    curNode.data = word


def solution(words, queries):
    answer = []
    head = Node('')
    bhead = Node('')
    cnt_dict = {}
    for word in words:
        insert(word, head)
        insert(word[::-1], bhead)
        if len(word) in cnt_dict:
            cnt_dict[len(word)] += 1
        else:
            cnt_dict[len(word)] = 1

    for query in queries:
        if query[0] == '?' and query[-1] == '?':  # ??????일떄
            try:
                answer.append(cnt_dict[len(query)])
            except:
                answer.append(0)
        else:
            org_len = len(query)
            if query[0] == '?':
                query = query[::-1]
                cur = bhead
            else:
                cur = head
            query = query.split("?")[0]
            idx = 0
            flag = 0
            for char in query:
                if char in cur.children:
                    cur = cur.children[char]
                else:
                    flag = 1
                    break
            if flag:
                answer.append(0)
            else:
                try:
                    answer.append(cur.length[org_len])
                except:
                    answer.append(0)

    return answer