[백준/python] 5670 휴대폰 자판

import sys; input = sys.stdin.readline


def insert(word):
    now = trie
    for char in word:
        now = now.setdefault(char, dict())
    now[0] = True

def search(now, cnt=0):
    global ans
    
    for char in now:
        if char == 0:
            ans += cnt
            continue

        if len(now)==1:
           search(now[char], cnt)
           continue

        search(now[char], cnt+1)    

while True:
    try:
        N = int(input())
        words = [input().strip() for _ in range(N)]
        
        trie = dict()
        for word in words:
            insert(word)

        ans = 0
        search(trie)

        if len(trie) == 1:
            ans += N
        
        print(f"{ans/N:.2f}")

    except:
        break

트라이

간단히 딕셔너리로 트라이구조를 구현했습니다.

그 후, search 함수를 통해 해당 위치의 trie의 원소가 1개만 있는 경우 카운트를 올리지 않고 내려가고, 그 이상인 경우 카운트를 올려 내려가는 방식으로 구현했습니다.

그리고,

3
h
he
hi

와 같이 첫 철자가 하나인 경우를 처리하기 위해 trie의 원소 개수가 1인 경우, 해당 N만큼 더해 예외 처리를 수행했습니다.