[BOJ 14725] 개미굴 (Python)

풀이

처음에 주어진 먹이 정보를 오름차순으로 정렬한 후 주어진 먹이 정보의 맨 왼쪽의 원소가 겹치는 것은 출력하지 않는 방식으로 접근하였다. O(N * K) -> 1,000 * 15 = 15,000의 시간 복잡도를 가지기 때문에 올바른 방식이라고 생각하였다. 하지만 출력 초과가 발생했다. 발생한 이유는 맨 왼쪽의 원소(루트 노드)의 중복 여부만 확인할 뿐 자식 노드의 중복 여부는 확인하지 않기 때문이었다.

그래서 자식 노드의 중복 여부도 확인하는 코드를 구현하였다. 그리고 주의할 점은 각 리스트의 맨 왼쪽(첫 번째) 원소부터 연속해서 겹치는 것이 있는지 확인해야 한다. 예를 들어

[A, B, C, D, E] // [A, B, C, Y, Z] 2개 리스트는 첫 번째 원소부터 연속해서 겹치는 것은 A, B, C이다. 따라서 idx=3이다.

하지만, [F, B, C, D, E] // [A, B, C, Y, Z] 2개 리스트는 첫 번째 원소부터 연속해서 겹치는 것은 하나도 없다. 따라서 idx=0이다.

추가로 이 문제는 문자열 트리 형태의 트라이 자료구조를 활용하여 푸는 방법도 있다. (트라이를 활용하는 것은 아직 미숙하다. 익숙해지도록 노력해야겠다...)

코드

정렬 후 완전탐색

import sys

input = sys.stdin.readline
N = int(input())
food_info = []
for i in range(N):
    data = list(input().split())
    food_info.append(data[1:])

food_info.sort()

dash = '--'
answer = []
for i in range(N):
    # 첫 번째는 부모, 자식 노드의 중복이 없으므로 그대로 출력.
    if i == 0:
        for j in range(len(food_info[i])):
            answer.append(dash * j + food_info[i][j])
    else:
        idx = 0   # 이전의 리스트와 현재 리스트에서 맨 왼쪽부터 겹치는 원소의 개수를 저장.
        for j in range(len(food_info[i])):
            # 이전의 리스트의 원소가 없거나, 이전의 리스트와 현재 리스트가 겹치지 않을 때
            if food_info[i - 1][j] != food_info[i][j] or len(food_info[i - 1]) <= j:
                break
            # 겹치는 원소가 존재한다면 해당 원소를 출력할 필요가 없으므로 idx를 조정.
            else:
                idx = j + 1
        for j in range(idx, len(food_info[i])):
            answer.append(dash * j + food_info[i][j])

for ans in answer:
    print(ans)

트라이 사용

# 참고 : https://cotak.tistory.com/3
import sys


class Trie:
    def __init__(self):
        self.root = {}

    def add(self, foods):
        cur = self.root

        for food in foods:
            if food not in cur:
                cur[food] = {}  # 자식 노드
            cur = cur[food]
        cur[0] = True  # 리프 노드 표시
        # print(self.root)
	
    # DFS
    def travel(self, level, cur):
        if 0 in cur:
            return
        cur_child = sorted(cur)

        for ch in cur_child:
            print("--" * level + ch)
            self.travel(level + 1, cur[ch])


input = sys.stdin.readline
N = int(input())
trie = Trie()
for i in range(N):
    data = list(input().split())
    trie.add(data[1:])

trie.travel(0, trie.root)