[백준 14725번] 파이썬 - 개미굴

백준 14725_개미굴

시간 제한 1초, 메모리 제한 256MB

# 조건

• 행복의 비결이 개미가 사는 개미굴에 있다고 생각한 윤수는 개미굴의 구조를 알아보기 위해 로봇 개미를 만들었다.
• 로봇 개미는 센서가 있어 개미굴의 각 층에 먹이가 있는 방을 따라 내려가다 더 이상 내려갈 수 없으면 그 자리에서 움직이지 않고 신호를 보낸다.
• 이 신호로 로봇 개미는 개미굴 각 층을 따라 내려오면서 알게 된 각 방에 저장된 먹이 정보를 윤수한테 알려줄 수 있다.

• 로봇 개미 개발을 완료한 윤수는 개미굴 탐사를 앞두고 로봇 개미를 테스트 해보기 위해 위 그림의 개미굴에 로봇 개미를 투입했다.
• 로봇 개미의 수는 각 개미굴의 저장소를 모두 확인할 수 있을 만큼 넣는다.
• 다음은 로봇 개미들이 윤수에게 보내준 정보다.
  • KIWI BANANA
  • KIWI APPLE
  • APPLE APPLE
  • APPLE BANANA KIWI
• 공백을 기준으로 왼쪽부터 순서대로 로봇 개미가 각 층마다 지나온 방에 있는 먹이 이름을 뜻한다.
• 윤수는 로봇 개미들이 보내준 정보를 바탕으로 다음과 같이 개미굴의 구조를 손으로 그려봤다.

APPLE
--APPLE
--BANANA
----KIWI
KIWI
--APPLE
--BANANA

• 개미굴의 각 층은 "--" 로 구분을 하였다.
• 또 같은 층에 여러 개의 방이 있을 때에는 사전 순서가 앞서는 먹이 정보가 먼저 나온다.
• 우리의 천재 공학자 윤수는 복잡한 개미굴들을 일일이 손으로 그리기 힘들어 우리에게 그려달라고 부탁했다.
• 한치 앞도 모르는 험한 이세상 그렇지만 오늘도 행복한 개미들!
• 행복의 비결을 알기 위해 윤수를 도와 개미굴이 어떤 구조인지 확인해보자.

입력

• 첫 번째 줄은 로봇 개미가 각 층을 따라 내려오면서 알게 된 먹이의 정보 개수 N (1 ≤ N ≤ 1000)개가 주어진다.
• 두 번째 줄부터 N+1 번째 줄까지, 각 줄의 시작은 로봇 개미 한마리가 보내준 먹이 정보 개수 K (1 ≤ K ≤ 15)가 주어진다.
• 다음 K개의 입력은 로봇 개미가 왼쪽부터 순서대로 각 층마다 지나온 방에 있는 먹이 정보이며 먹이 이름 길이 t는 1 ≤ t ≤ 15를 만족한다.
• 먹이 정보는 알파벳 대문자로만 이루어져 있다.

출력

• 개미굴의 시각화된 구조를 출력하여라.
• 개미굴의 각 층을 "--" 로 구분하며, 같은 층에 여러개의 방이 있을 때에는 사전 순서가 앞서는 먹이 정보가 먼저 나온다.
• 최상위 굴을 포함하여 하나의 굴에서 개미굴이 여러개로 나뉠 때 먹이 종류별로 최대 한 번만 나올 수 있다.

# 접근 방법

• 처음엔 최상위 굴에 있는 먹이 이름을 key로 사용하여 딕셔너리를 활용해주었다.
  • value에 새로운 먹이가 들어올 때마다 저장 해주고, 사전순 정렬을 하여 -- 와 함께 출력하였다.
  • 제출 결과는 출력초과..
• 해당 문제의 조건에 하나의 굴에서 개미굴이 여러개로 나뉠 때 먹이 종류별로 최대 한 번만 나올 수 있다는 것 때문에 불필요한 출력이 나와서 그런 것 같았다.

import sys  
sys.stdin = open('input.txt')  
input = sys.stdin.readline  
from collections import defaultdict  

N = int(input())  
result = defaultdict(list)  
for _ in range(N):  
    query = [*map(str, input().split())]  
    struct = query[1:]  
    result[struct[0]].append(struct[1:])  

res = dict(sorted(result.items()))  
for i in res:  
    j = sorted(res[i])  
    print(i)  
    flag = len(j)  
    idx = 0  
    while idx < flag:  
        slash = 1  
        for k in j[idx]:  
            print('--'*slash + k)  
            slash += 1  
        idx += 1

• 따라서, 트라이 구조를 딕셔너리로 구현해주었다.
• query 한 줄을 받을 때마다 can_add 함수를 시작해주었다.
  • 먹이의 제일 첫 번째 값이 이미 result 딕셔너리에 존재한다면 pass, 존재하지 않는다면 한 depth 들어가서 A : { B : { C : {}}} 와 같이 key 값으로 추가해주었다.
• 출력도 마찬가지로, key값을 기준으로 정렬해주고, -- 를 늘려가며 출력해준다.

import sys  
sys.stdin = open('input.txt')  

def can_add(val, arr):  
    if len(arr) == 0:  
        return  

    if arr[0] not in val:  
        val[arr[0]] = {}  
    can_add(val[arr[0]], arr[1:])  


def printTree(dic, leng):  
    for i in sorted(dic.keys()):  
        print("--"*leng+i)  
        printTree(dic[i],leng+1)  

N = int(input())  
query = [list(map(str, input().split())) for _ in range(N)]  
trie = {}  

for i in query:  
    i = i[1:]  
    can_add(trie, i)  

print(trie)  
printTree(trie, 0)