[Python] BOJ 14725 - 개미굴

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

14725번: 개미굴

# Trie(트라이) 자료구조로 푸는 문제이다.

 

1. Trie클래스를 선언한다.

   1-1. insert메소드: 입력받은 먹이들(food배열)을 반복문으로 돌면서 cur_node의 key에 먹이가 없다면

               cur_node[현재먹이] = {}로 만들어주고, cur_node의 key에 먹이가 존재한다면

               cur_node = cur_node[string]으로 다음 먹이를 가리키도록 한다.

 

               먹이의 마지막에는 '*'을 삽입해준다.

 

   1-2. print_list메소드: 마지막에 개미굴의 시각화된 구조를 출력하기 위한 메소드이다.

               처음에 cur_node에 trie.root가 들어온다. 만약 '*' in cur_node이면 먹이의 마지막이므로

                return을 해준다.

 

                cur_node의 key값을 반복문으로 돌며 cnt(현재깊이) * '--'  + child(현재 먹이) 를 출력하고

                self.print_list(cnt+1, cur_node[child]) 로 재귀호출을 한다. 

                즉, 먹이의 끝에 도달할때까지 계속 dfs형식으로 출력해주는 것이다.

 

2. N만큼 k와 먹이정보 t를 입력받고 먹이들을 food_list에 append해준다.

    "같은 층에 여러개의 방이 있을 때에는 사전 순서가 앞서는 먹이 정보가 먼저 나온다."

      모든 정보 N개를 입력받은 뒤  위의 조건을 만족시키기 위해 food_list를 sort해준다.

 

3. 정렬한 뒤 food_list에 있는 먹이들을 trie에 insert해준다.

    trie.print_list(0, trie.root)를 해주고 맨 위의 개미굴 입구로부터 개미굴 내부구조를 시각화해서 출력한다.

 

import sys
input = sys.stdin.readline
class Trie:
    def __init__(self):
        self.root = {}      # 자식저장
    
    def insert(self, food):
        cur_node = self.root

        for string in food:
            if string not in cur_node:
                cur_node[string] = {}
            cur_node = cur_node[string]

        cur_node['*'] = True       # 문자열 마지막에 * 삽입
    
    def print_list(self, cnt, cur_node):
        if '*' in cur_node:
            return

        for child in cur_node:
            print("--" * cnt + child)
            self.print_list(cnt+1, cur_node[child])

trie = Trie()
N = int(input())
food_list = []
for _ in range(N):
    k, *t = map(str,input().split())
    food_list.append(t)

food_list.sort()
for foods in food_list:
    trie.insert(foods)

trie.print_list(0, trie.root)

(solve.ac 티어: 골드 2)