Baekjoon Online Judge

algorithm practice

- 단계별 문제풀기

34. 문자열 알고리즘 1

KMP 알고리즘과 트라이 자료구조를 다뤄 봅시다.

5. 개미굴

14725번

트라이에 대한 감을 잡는 문제

이번 문제는 로봇 개미들이 보내준 문자열 바탕으로 개미굴이 어떤 구조인지 확인하는 문제이다.
로봇 개미가 각 층을 따라 내려오면서 알게 된 먹이의 정보 개수 N개(1 ≤ N ≤ 1000),
N+1 번째 줄까지, 각 줄의 시작은 로봇 개미 한마리가 보내준 먹이 정보 개수 K(1 ≤ K ≤ 15), 다음 K개의 입력은 로봇 개미가 왼쪽부터 순서대로 각 층마다 지나온 방에 있는 먹이 정보이며 먹이 이름 길이 t는 (1 ≤ t ≤ 15) 이다.

트리 구조의 이해와 dfs 활용이 문제의 핵심으로 생각된다.

Java / Python

• Java

import java.io.*;
import java.util.*;

public class Main {
	static StringBuilder sb;

	public static class Trie {
		// 이름, 자식 저장
		ArrayList<Trie> list;
		String name;

		Trie(String name) {
			list = new ArrayList<>();
			this.name = name;
		}
	}

	public static void main(String[] args) throws Exception {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
		StringTokenizer st;
		sb = new StringBuilder();

		int N = Integer.parseInt(br.readLine());

		Trie trie = new Trie("");
		Trie node;

		while (N-- > 0) {
			st = new StringTokenizer(br.readLine());
			int k = Integer.parseInt(st.nextToken());
			node = trie;

			while (k-- > 0) {
				String name = st.nextToken();
				int idx = -1;

				for (int i = 0; i < node.list.size(); i++) {
					if (node.list.get(i).name.equals(name)) {
						idx = i;
						break;
					}
				}

				if (idx == -1) {
					// 현재 노드의 list에 name 추
					// 노드 추가한 변수의 위치로 이동
					node.list.add(new Trie(name));
					node = node.list.get(node.list.size() - 1);
				} else {
					// 자식으로 이
					node = node.list.get(idx);
				}
			}
		}

		print(trie, -1);
		bw.write(sb.toString());
		bw.flush();
		bw.close();
		br.close();
	}

	public static void print(Trie trie, int depth) {
		// 재귀 형식
		// 현재 노드의 list를 이름 순으로 정렬하고 현재 노드의 name을 StringBuilder에 추가
		// list의 값에 차례대로 접근하여 dfs
		Collections.sort(trie.list, new Comparator<Trie>() {
			@Override
			public int compare(Trie o1, Trie o2) {
				return o1.name.compareTo(o2.name);
			}
		});
		if (depth != -1) {
			for (int i = 0; i < depth; i++) {
				sb.append("--");
			}
			sb.append(trie.name).append("\n");
		}

		for (Trie tr : trie.list) {
			print(tr, depth + 1);
		}
	}
}

• Python

import sys
input = sys.stdin.readline

N = int(input())
ant = {}

for i in range(N):
    name = list(input().split())
    target_dict = ant
    for j in name[1:]:
        if j not in target_dict:
            target_dict[j] = {}
        target_dict = target_dict[j]

def getResult(t, i):
    target_key = sorted(t.keys())
    for s in target_key :
        print('--'*i + s)
        getResult(t[s],i+1)

getResult(ant,0)