백준 5670번: 휴대폰 자판 (Python, PyPy3)

접근

트라이를 이용하여 풀었다. 트라이를 생성하여 입력받는 문자열들을 트라이에 저장해주고, 나중에 평균을 구해주기 위해 문자열을 저장할 때, 해당 string 이 몇번 사용되었는지 각 노드에 저장하는 식을 추가해주었다.

평균을 구할 때에는 해당 노드의 자식 수에 따라서 자식이 한개일 경우 카운트하지 않고, 자식이 있을 경우 각 자식이 몇번 쓰이는지를 체크하여 모두 더해주었다. 더해진 모든 수를 총 단어 수로 나누면 평균 입력 수를 구할 수 있었다.

Python 으로는 50% 시간 초과 발생하여 PyPy3으로 문제 해결하였다.

트라이 기본 문제: 2021.06.19 - [코딩/백준 (Python)] - 백준 14725번: 개미굴 (Python)

코드

import sys
from collections import deque


class Node:
    def __init__(self, key):
        self.key = key
        self.children = dict()
        self.cnt = 0


class Trie:
    def __init__(self):
        self.head = Node(None)

    def insert(self, string):
        curr_node = self.head
        curr_node.cnt += 1
        for char in string:
            if char not in curr_node.children:
                curr_node.children[char] = Node(char)
            curr_node = curr_node.children[char]
            curr_node.cnt += 1

    def average(self):
        curr_node = self.head
        count = self.head.cnt
        q = deque([])
        for child in curr_node.children:
            q.append(curr_node.children[child])
        while q:
            curr = q.popleft()
            if len(curr.children) > 1:
                for child in curr.children:
                    if child != "\n":
                        q.append(curr.children[child])
                        count += curr.children[child].cnt
            else:
                for child in curr.children:
                    q.append(curr.children[child])
        return count / self.head.cnt


while 1:
    trie = Trie()
    try:
        n = int(sys.stdin.readline())
    except:
        break
    for __ in range(int(n)):
        trie.insert(sys.stdin.readline())
    print("%.2f" % trie.average())

더 생각해 볼 것?

이번 코드는 조금 지저분해 보여서 문제를 좀 더 깔끔하게 풀 수 있는 알고리즘이 있는지 고민을 좀 더 해봐야 겠다. 그리고 더 개선된 알고리즘으로 풀었을 때 시간이 더 단축될 수 있는지도 확인해보아야겠다.

코드나 내용 관련 조언, 부족한 점 및 질문 언제든 환영합니다!