[프로그래머스] [3차] 자동완성 - Trie - Lv.4

 

https://school.programmers.co.kr/learn/courses/30/lessons/17685

프로그래머스

 

문제 분석

어디까지 입력해야 자동완성이 내가 찾으려는 단어만 딱 표시하는지를 찾으면 되는 문제다.

 

자료구조 (Trie)

단어는 맨 첫글자부터 순차적으로 입력된다는 것을 생각했을 때, 일종의 트리 구조처럼 각 글자를 따라 내려가면 되겠다는 생각을 했다.

 

이럴 때 사용하는 자료구조가 바로 Trie다.

문자열 트리구조인데, 예시를 보면 간단하다.

 

예시

go
gone
guild

 

아래는 위 단어들을 가지고 트리를 만든 구조이다.

각 노드는 아래같은 구조로 만들면 공간 낭비 덜하고 쉽게 만들 수 있다.

class Node:
    def __init__(self, letter):
        self.children = [None] * 26

 

본인의 정보는 딱히 필요 없다.

탑다운으로 탐색되고, 부모노드가 자식노드의 정보를 알고있기 때문이다.

 

문제 풀이

1. 학습

문제는 주어진 단어를 모두 학습했다는 가정하에 풀게 되어있다.

즉, 주어진 단어를 가지고 미리 트리를 만들어야한다는 것이다.

 

따라서 노드에 자식 노드를 추가하는 메소드를 구현했다.

class Node():
    def __init__(self):
        self.count = 1
        self.children = [None] * 26
    
    def add(self, child):
        child_ord = ord(child) - ord('a')
        if self.children[child_ord] is None:
            self.children[child_ord] = Node()
        else:
            self.children[child_ord].count += 1

여기서 표시되는 count는 추후에 설명하겠다.

 

1. ord를 통해 추가할 문자의 ascii 값을 구한다.

2. 자식노드 배열의 해당 문자 위치에 새 Node를 추가한다.

 

헤드에 다음 노드를 붙여야하기 때문에 아래 메소드도 만들어줬다.

    def get_next_node(self, next_letter):
        next_letter_ord = ord(next_letter) - ord('a')
        return self.children[next_letter_ord]

 

이후 루트 노드를 만들고, head를 하나 만들어서 노드를 붙이도록 하면 트리가 완성된다.

    trie = Node()
    for word in words:
        head = trie
        for letter in word:
            head.add(letter)
            # 다음 letter 삽입 가능하도록 맵 구성
            head = head.get_next_node(letter)

 

1. 각 단어의

2. 각 문자를 순회하면서 노드를 붙인다.

 

2. 자동완성 언제까지?

먼저 문제의 일부분을 발췌했다.

go를 찾을 때 go를 모두 입력해야 한다.
gone을 찾을 때 gon 까지 입력해야 한다. (gon이 입력되기 전까지는 go 인지 gone인지 확신할 수 없다.)
guild를 찾을 때는 gu 까지만 입력하면 guild가 완성된다.

 

이제 핵심부분이다. 언제까지 문자를 타이핑 해야하는지 체크해야한다.

여러 삽질을 했었는데, 찾아낸 결론은 다음과 같다.

 

1. Node에 count를 추가한다.

2. 여러 단어의 앞에서 n번째까지의 문자열이 겹칠 경우, 그 노드들은 count에 값을 1씩 더한다.

 

즉, 트리 순서에서 해당 문자가 나온 횟수를 체크한다.

이렇게 하면 몇개의 단어가 어디까지 문자열이 겹치는지 쉽게 확인할 수 있다.

 

예시

위 트리에서는

g는 3개의 단어가 겹치고 (go, gone, guild)

go는 2개의 단어가 겹치고 (go, gone)

gui는 1개의 단어가 겹치고 (guild)

gon은 1개의 단어가 겹치는 것을 알 수 있다. (gone)

 

따라서 go를 검색한다 했을 때,

go까지 검색해도 go와 gone가 자동완성으로 뜬다는 것이다.

=> 글자 2개 입력해야함

 

gone를 검색한다 하면,

go까지는 go와 gone가 뜨지만

gon까지 검색하면 gone만 자동완성으로 뜬다.

=> 글자 3개만 입력하면 됨

 

코드를 완성하면 다음과 같다.

    sum_counts = 0
    for word in words:
        head = trie
        count = 0
        for letter in word:
            head = head.get_next_node(letter)
            count += 1
            if head.count == 1:
                break
      
        sum_counts += count
    
    return sum_counts

 

count를 계속 늘리며 트리를 내려가다가

더이상 문자열이 겹치는 단어가 없으면 탐색을 종료한다.

 

 

최종 코드

def solution(words):
    trie = Node(None)
    for word in words:
        head = trie
        for letter in word:
            head.add(letter)
            # 다음 letter 삽입 가능하도록 맵 구성
            head = head.get_next_node(letter)
    
    # 문자열 문자 하나씩 타고 내려가는데, 더이상 본인 외의 자식이 없다면 
    sum_counts = 0
    for word in words:
        head = trie
        count = 0
        for letter in word:
            # 해당 문자 위치에 대한 노드 가지고 옴
            head = head.get_next_node(letter)
            count += 1
            if head.count == 1:
                break
      
        sum_counts += count
    
    return sum_counts
            
    
class Node():
    def __init__(self, letter):
        self.letter = letter
        self.count = 1
        self.children = [None] * 26
    
    def add(self, child):
        child_ord = ord(child) - ord('a')
        if self.children[child_ord] is None:
            self.children[child_ord] = Node(child)
        else:
            self.children[child_ord].count += 1
    
    def get_next_node(self, next_letter):
        next_letter_ord = ord(next_letter) - ord('a')
        return self.children[next_letter_ord]

 

 

 

후기

트리 만들기는 쉬웠다.

그런데 어떻게 탐색하면 좋을지에서 헤매었다.

stack도 생각해보고, 단순히 내려가면서 자식이 없으면 빠져나오는 것도 구상해봤는데, 손으로 구상하는 단계에서부터 예외에 걸렸다.

 

그래도 count를 생각해내면서부터는 쉽게 해결할 수 있었다.