[Python] 백준 14725번: 개미굴

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

14725번: 개미굴

 

import sys
input = sys.stdin.readline


class Node():
    def __init__(self, key):
        self.key = key
        self.children = {}


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

    def insert(self, strings):
        current_node = self.head

        for string in strings:
            # 없다면 생성
            if string not in current_node.children:
                current_node.children[string] = Node(string)

            current_node = current_node.children[string]

    def dfs(self, current_node, depth):
        for k, v in sorted(current_node.children.items()):
            # 출력
            for _ in range(depth):
                print("--", end="")
            print(k)

            # 재귀
            self.dfs(v, depth + 1)


if __name__ == '__main__':
    N = int(input())

    trie = Trie()
    for _ in range(N):
        cur_input = list(map(str, input().split()))
        trie.insert(cur_input[1:])

    trie.dfs(trie.head, 0)

 

트라이 알고리즘 기초 문제이다.

 

트라이 알고리즘은, 트리를 이용한 알고리즘이고, 이 문제를 위한 코드(이 문제는 일반적인 Trie보다 간결하고 쉬운 코드를 요구한다.)에서는 각 노드가 key, children을 가진다.

이 문제에서 설명할 때 제시한 그림에서, 1층에 있는 APPLE을 기준으로 설명하겠다. key는 해당 노드의 데이터로, APPLE이 이것에 해당한다. children은 하위에 있는 APPLE, BANANA를 담는다.

 

insert() 작업은 입력받은 리스트를 for문으로 돌면서, 자식 노드에 현재 문자열이 없으면 children에 노드를 생성하고, 계속해서 탐색을 이어나간다.

 

dfs() 함수는, 이 문제가 출력을 dfs 구조로 하기 때문에 짠 것이며, 말그대로 DFS로 탐색을 하고, 현재 재귀 함수의 깊이를 통해 출력을 해주는 식이다.