백준 5670 <휴대폰 자판> Python

https://www.acmicpc.net/problem/5670

5670번: 휴대폰 자판

---

트라이 자료구조를 사용해 풀 수 있는 문제였습니다. 처음에는 현재 노드의 자식이 하나면 카운트를 스킵했지만 어떤 단어의 부분집합이 함께 사전에 포함되어 있는 경우, 예를 들어 hello 와 hell 이 같은 사전에 포함되어 있는 경우, hello 의 버튼 클릭 횟수를 카운팅 할 때 o 를 자동으로 누르게 되어 카운팅이 제대로 되지 않습니다. 그러므로 자식의 수 + (value 가 존재할 때 1 / value 가 존재하지 않을 때 0) 을 해주어 해당 단어가 포함되어 있을 경우를 고려해야 합니다.

또 해당 문제는 테스트 케이스의 수가 주어지지 않으므로 EOF 를 고려해야합니다. 

---

import sys

def solution(N, word_list):

    # 단어들의 버튼 클릭 횟수 총 합

    btn_cnt = 0

    # 트라이

    trie = Trie()

    # 단어 입력

    for word in word_list:

        # Trie 의 insert 메서드

        trie.insert(word)

    # 단어 버튼 클릿 횟수 카운팅

    for word in word_list:

        # Trie 의 count 메서드

        btn_cnt += trie.count(word)

    # 평균값 출력

    print(f'{round(btn_cnt/N, 2):.2f}')

# 트라이에 들어갈 노드 클래스

class Node:

    # 초기화

    def __init__(self, key, value=None):

        # children 의 key-value 쌍에서의 key

        self.key = key

        # children 의 key-value 쌍에서의 value

        self.value = value

        # children

        self.children = {}

# 트라이 자료구조

class Trie:

    # 초기화

    def __init__(self):

        # 최상단 노드

        self.head = Node(None)

   

    # 입력 함수

    def insert(self, string):

        # 현재 비교하는 문자 노드

        current_node = self.head

        # 순회

        for char in string:

            # char 이 현재 노드의 자식에 속해잇지 않으면 삽입

            if char not in current_node.children:

                current_node.children[char] = Node(char)

            # 노드 이동

            current_node = current_node.children[char]

        # 문자열을 모두 순회했으면 끝에 value 추가

        current_node.value = string

   

    # 버튼 클릭 횟수 카운트 함수

    # 문제에 맞게 변형

    def count(self, string):

        # 버튼 클릿 횟수

        btn_cnt = 0

        # 현재 비교하는 문자 노드

        current_node = self.head

        # 순회

        for char in string:

            # 첫번째 문자는 추론하지 않음

            if not current_node.key:

                # 노드 이동

                current_node = current_node.children[char]

                # 버튼 클릭

                btn_cnt += 1

            # 두번째 문자부터

            else:

                # 자식의 수와 value의 수의 합이 1 이면

                if len(current_node.children) + bool(current_node.value) == 1:

                    # 노드 이동

                    current_node = current_node.children[char]

                # 자식이 하나가 아니면

                else:

                    # 노드 이동

                    current_node = current_node.children[char]

                    # 버튼 클릭

                    btn_cnt += 1

        return btn_cnt

    # starts_with 함수는 문제 해결에 필요없어서 구현하지 않았음

       

while True:

    try:

        N = int(sys.stdin.readline())

        word_list = [sys.stdin.readline().strip() for _ in range(N)]

        solution(N, word_list)

    except:

        break