프로그래머스 / 2018 KAKAO 자동완성 / 파이썬 / 트라이

programmers.co.kr/learn/courses/30/lessons/17685

Trie 자료구조를 안다면 쉬운 문제이고 모르면 어려운 문제입니다.

Trie 자료구조를 조금 변형하여 쉽게 구현할 수 있었습니다.
Trie에 cnt를 추가해주면 쉽게 풀 수 있습니다.

cnt는 같은 루트를 공유하는 경우를 나타냅니다.
cnt가 1이라면 더 이상 찾아볼 이유가 없습니다.

그리고 cnt는 문자열 길이를 넘을 수 없음에 유의합니다.

class Trie:
    def __init__(self):
        self.root = Node()

    def insert(self, s):
        node = self.root
        for c in s:
            if c in node.child:
                node.child[c].cnt += 1
            else:
                node.child[c] = Node()
            node = node.child[c]
        node.child['*'] = Node()

    def search(self, s):
        node = self.root
        cnt = 1
        for c in s:
            node = node.child[c]
            if node.cnt == 1:
                return cnt
            cnt += 1
        return min(cnt, len(s))

class Node:
    def __init__(self):
        self.child = {}
        self.cnt = 1

def solution(words):
    answer = 0
    root = Trie()
    for word in words:
        root.insert(word)
    for word in words:
        answer += root.search(word)
    return answer