백준 5670번 파이썬 문제 풀이

문제

휴대폰에서 길이가 P인 영단어를 입력하려면 버튼을 P번 눌러야 한다. 그러나 시스템프로그래밍 연구실에 근무하는 승혁연구원은 사전을 사용해 이 입력을 더 빨리 할 수 있는 자판 모듈을 개발하였다. 이 모듈은 사전 내에서 가능한 다음 글자가 하나뿐이라면 그 글자를 버튼 입력 없이 자동으로 입력해 준다! 자세한 작동 과정을 설명하자면 다음과 같다.

1. 모듈이 단어의 첫 번째 글자를 추론하지는 않는다. 즉, 사전의 모든 단어가 같은 알파벳으로 시작하더라도 반드시 첫 글자는 사용자가 버튼을 눌러 입력해야 한다.
2. 길이가 1 이상인 문자열 c1c2...cn이 지금까지 입력되었을 때, 사전 안의 모든 c1c2...cn으로 시작하는 단어가 c1c2...cnc로도 시작하는 글자 c가 존재한다면 모듈은 사용자의 버튼 입력 없이도 자동으로 c를 입력해 준다. 그렇지 않다면 사용자의 입력을 기다린다.

예를 들면, 사전에 "hello", "hell", "heaven", "goodbye" 4개의 단어가 있고 사용자가 "h"를 입력하면 모듈은 자동으로 "e"를 입력해 준다. 사전상의 "h"로 시작하는 단어들은 모두 그 뒤에 "e"가 오기 때문이다. 그러나 단어들 중 "hel"로 시작하는 것도, "hea"로 시작하는 것도 있기 때문에 여기서 모듈은 사용자의 입력을 기다린다. 이어서 사용자가 "l"을 입력하면 모듈은 "l"을 자동으로 입력해 준다. 그러나 여기서 끝나는 단어인 "hell"과 그렇지 않은 단어인 "hello"가 공존하므로 모듈은 다시 입력을 기다린다. 사용자가 "hell"을 입력하기 원한다면 여기서 입력을 멈출 것이고, "hello"를 입력하기 원한다면 직접 "o" 버튼을 눌러서 입력해 줘야 한다. 따라서 "hello"를 입력하려면 사용자는 총 3번 버튼을 눌러야 하고, "hell", "heaven"은 2번이다. "heaven"의 경우 "he"에서 "a"를 입력하면 바로 뒷부분이 모두 자동으로 입력되기 때문이다. 비슷하게, "goodbye"는 단 한 번만 버튼을 눌러도 입력이 완료될 것이다. "g"를 입력하는 순간 뒤에 오는 글자가 항상 유일하므로 끝까지 자동으로 입력되기 때문이다. 이때 사전에 있는 단어들을 입력하기 위해 버튼을 눌러야 하는 횟수의 평균은 (3 + 2 + 2 + 1)/4 = 2.00이다.

사전이 주어졌을 때, 이 모듈을 사용하면서 이와 같이 각 단어를 입력하기 위해 버튼을 눌러야 하는 횟수의 평균을 구하는 프로그램을 작성하시오.

입력

입력은 여러 개의 테스트 케이스로 이루어져 있다.

각 테스트 케이스의 첫째 줄에 사전에 속한 단어의 개수 N이 주어지며(1 ≤ N ≤ 105), 이어서 N개의 줄에 1~80글자인 영어 소문자로만 이루어진 단어가 하나씩 주어진다. 이 단어들로 사전이 구성되어 있으며, 똑같은 단어는 두 번 주어지지 않는다. 각 테스트 케이스마다 입력으로 주어지는 단어의 길이 총합은 최대 106이다.

출력

각 테스트 케이스마다 한 줄에 걸쳐 문제의 정답을 소수점 둘째 자리까지 반올림하여 출력한다.

문제풀이

이번 문제 풀이에서는 트라이 구조를 사용했다. 문제를 풀 때 중요한 점은 트라이 구조 내에서 어떻게 자판을 입력하는 것을 카운트 해야하는지를 고려해야 했다. 그것만 잘 고려한다면, 문제는 이미 다 풀린 것이나 다름 없었다. 나는 다음과 같은 경우에 자판을 한번씩 더 눌러야 한다고 판단했다.

• 처음 문자를 칠 때

 :  이는 문제 설명에 명시되어있다. "모듈이 단어의 첫 번째 글자를 추론하지는 않는다."

• 현재 노드 기준 자식 노드가 2개 이상 있을 때

 : 자식 노드가 2개 이상 있다는 말은 현재 노드에서 분기가 갈린다는 말이다. 따라서 이 때는 자판을 직접 쳐주어야 한다.

• 탐색 문자 기준으로는 아직 더 탐색해야하지만, 현재 노드가 다른 문자 기준으로 마지막 지점일 때

: 예를 들어 HELLO를 검색할 때, HELL 상태에서 자식 노드는 'O' 하나 밖에 없지만, 'HELL'이 이미 완결된 상태이기 때문에 'HELLO'를 선택하기 위해선 'O'를 한번 더 쳐주어야 한다.

위와 같은 로직으로 문제를 풀기 위해서는 일전에 사용했던 노드 구조에 한 가지 변수가 더 추가되어야 한다. 현재 노드 기준으로 문자가 완성된 적이 있는지를 기록하고 확인하기 위한 Flag가 필요하다. 따라서 나는 노드 클래스를 아래와 같이 작성했다.

class Node(object) :
    def __init__(self,data):
        self.data = data
        self.child = {}
        self.flag = False

트라이 클래스는 초기화, 문자열 추가, 버튼 Count 함수를 만들었다. 문자열을 추가할 때에는 For문으로 문자열을 다 작성하고 난 다음, 마지막 노드의 Flag를 True로 바꾸어 주도록 해서, 여기가 어떤 문자의 마지막이라는 것을 표현했다. 

Count 함수는 해당 문자를 따라 노드를 지나가면서, 위의 로직을 만족할 때마다 Return_Value의 값을 하나씩 늘리는 것으로 작성했다. 

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

    def insert(self,char):
        cur = self.head

        for c in char :
            if c not in cur.child :
                cur.child[c] = Node(c)
            cur = cur.child[c]
        cur.flag = True

    def count(self,char):
        cur = self.head
        return_value = 0

        for c in char :
            #현재 노드가 Head일 때, 1번의 경우
            if cur == self.head :
                return_value +=1
			#현재 노드의 자식 노드가 2개 이상 or 현재 노드에서 문자가 끝났었던 적이 있다면 (2+3번)
            elif len(cur.child) > 1 or cur.flag == True :
                return_value +=1

            cur = cur.child[c]

        return return_value

소수점 아래 2자리 반올림 표현은 이렇다 할 방법을 찾기 어려워, 나는 다음과 같이 표현을 했다.

• 현재 값을 String Type으로 바꾼다.
• 현재 값을 소수점을 기준으로 a,b라는 변수에 나눠담는다.
• b의 문자열 길이가 2보다 작을 경우, 2가 될 때가지 0을 뒷쪽에 추가해준다.
• a 문자열 + 마침표 + b 문자열 추가한다.

코드로 표현하면 아래와 같다.

    answer = str(round(answer,2))
    a,b = answer.split('.')

    while len(b) < 2 :
        b = b + '0'
    answer = (a + '.' + b)

이 코드들을 합치면, 내 코드는 아래와 같다.

import sys

class Node(object) :
    def __init__(self,data):
        self.data = data
        self.child = {}
        self.flag = False



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

    def insert(self,char):
        cur = self.head

        for c in char :
            if c not in cur.child :
                cur.child[c] = Node(c)
            cur = cur.child[c]
        cur.flag = True

    def count(self,char):
        cur = self.head
        return_value = 0

        for c in char :
            if cur == self.head :
                return_value +=1

            elif len(cur.child) > 1 or cur.flag == True :
                return_value +=1

            cur = cur.child[c]

        return return_value


while 1 :
    try :
        n = int(sys.stdin.readline().rstrip())
    except :
        exit()
    a = Trie()
    string_list = []
    answer = []
    for i in range(n) :
        my_str = str(sys.stdin.readline().rstrip())
        a.insert(my_str)
        answer.append(my_str)


    my_sum = 0

    for my_str in answer :
        my_sum += a.count(my_str)

    answer = my_sum / n
    answer = str(round(answer,2))
    a,b = answer.split('.')

    while len(b) < 2 :
        b = b + '0'
    answer = (a + '.' + b)
    print(answer)