[Python] 프로그래머스 / 2018 KAKAO BLIND RECRUITMENT / [3차] 자동완성

문제 링크

 

트라이(Trie) 자료구조

최근에야 검색창 자동완성 기능에 NLP가 활용되지만, 고전적으로는 트라이 자료구조가 활용된다. 여러 단어를 Tree 자료 구조로 관리하는 것이다. 이 문제에서 예시로 들고 있는, gone, guild를 예시로 들면 다음과 같다 :

       (root)
        |
        g
       / \
      o   u
     /     \
    n       i
   /         \
  e           l
              |
              d

 

더 자세한 설명은 이 블로그에서 상세히 작성해놓은 것이 있다.

 

 

정답 코드

#트라이 자료구조 만들기
class TrieNode:
    def __init__(self):
        self.children = {}
        self.count = 0

def insert_word(root, word):
    node = root
    for char in word:
        if char not in node.children:
            node.children[char] = TrieNode()
        node = node.children[char]
        node.count += 1

def find_prefix_length(root, word):
    node = root
    for i, char in enumerate(word):
        if node.children[char].count == 1:
            """
            count가 1이라는 것은, 이 분기점에서 char로 시작하는 단어가 유일한 경우라는 뜻

            예를 들어,
            apple
            ape
            apricot
            에서,

            p의 node.children = {'p' : TrieNode(), 'e' : TrieNode(), 'r' : TrieNode()}
            이고, 각 value의 count는 1이다.
            """
            return i + 1
            
        node = node.children[char]
    return len(word)

def solution(words):
    root = TrieNode()
    for word in words:
        insert_word(root, word)
    
    total_length = 0
    for word in words:
        total_length += find_prefix_length(root, word)
    
    return total_length