[PS] 백준 14725번 개미굴

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

문제설명

개미 굴이 있다. 개미 굴에 로봇 개미를 보내면 로봇 개미는 굴의 아래층으로만 향한다. 아래층으로 향하며 더이상 내려갈 층이 없으면, 방문한 개미 굴 방에있는 먹이의 이름을 방문 순서대로 데이터를 전송한다.(데이터는 가장왼쪽 방부터 오른쪽 방순으로 전송된다) 모든 개미 굴 방을 방문할 수 있는 수의 로봇 개미를 투입했을 때, 로봇 개미가 보내오는 데이터를 토대로, 트리 형태로 나타내면 된다.(사전순으로 정렬 된 순서로)

문제풀이

1. 트리 형태로 나타내면 된다.(라고 문제에 주어져있다) 개미 굴 그림을 보면, 먹이 이름(문자열)이 노드의 데이터가 되면 되고, 사전순으로 정렬 된 순서로 다음 노드(먹이 이름)를 방문해야 하므로 정렬된 상태로 저장하면서, 다음 노드로 이동할 수 있으면서, 해당 노드가 존재하는지 하지 않는지를 준수한 복잡도안에 계산할 수 있는 TreeMap을 사용하면 된다는 것을 생각해 낼 수 있다. 따라서 노드의 형태는 아래와 같다.

class Node { 
	String name;
    TreeMap<String, Node> childs;
}

우선 root Node를 만들어놓고, 차례대로 입력을 받으며 각 굴의 깊이에서 현재 먹이 이름이 현재 Node의 TreeMap에 Key로 존재하는지 검사하면 된다.

1) 존재하면 현재 노드를 그 Node로 변경하고 반복
2) 존재하지 않는다면 노드를 새로 생성하고 현재 노드의 TreeMap에 추가후 추가한 노드로 현재 Node를 변경하고 반복하면 된다.

이 문제 카테고리가 Trie로 되어있는데, Trie는 한 문자가 노드가 된다. 결론적으로 접두사를 트리형태로 만들어 검색속도를 대폭 증가시킬 수 있는(검색 문자 길이 := m, O(m)) 자료구조이다. 접두사를 하나씩 보며 일치하는 것을 찾아가므로 backtracking이 필요없다. 따라서 Trie를 트리형태인 결정적 유한 오토마타(DFA)로 볼 수도 있다고 한다.
그리고 Trie를 이용해서, 문자열 자동 완성을 할 수도 있다. 임의의 노드 x까지 지나온 문자들을 이어붙이면 접두사가 되는데, 그러면 노드 x 아래 레벨에 존재하는 모든 노드들은 동일한 접두사를 공유하기 때문이다. 자세한건 더 공부하도록 하자.
개미굴 문제에서는 각 노드에 문자열을 저장하는 형태가 된다.

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.Scanner;
import java.util.StringTokenizer;
import java.util.TreeMap;

class Main {
	static BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
	static BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
	static Scanner sc = new Scanner(System.in);

	static StringTokenizer st;
	
	static int n;
	static Node root;
	
	/**
	 * 트리의 노드
	 */
	static class Node {
		String name; // 먹이 이름
		TreeMap<String, Node> childs; // 연결된 아래층의 먹이들을 먹이 이름의 사전순으로 정렬하여 보관
		
		public Node(String name) {
			super();
			this.name = name;
			this.childs = new TreeMap<>();
		}
	}
	
	public static void main(String[] args) throws IOException {
		n = Integer.parseInt(br.readLine());
		Node root = new Node("root");
		
		
		for (int i = 0; i < n; i++) {
			Node curNode = root;
			st = new StringTokenizer(br.readLine());
			int k = Integer.parseInt(st.nextToken()); // 로봇 개미 한마리가 보내준 먹이 정보 개수 K
			
			for (int j = 0; j < k; j++) {
				String word = st.nextToken(); // 먹이 이름
				
				// 해당 굴, 해당 층에 먹이이름'word'가 이미 있는 경우 아래 층으로 이동
				if (curNode.childs.containsKey(word)) {
					curNode = curNode.childs.get(word);
					continue;
				}
				
				// 해당 굴, 해당 층에 먹이이름'word'가 없는 경우 먹이이름'word'를 생성하고 연결한 후 아래 층으로 이동
				curNode.childs.put(word, new Node(word));
				curNode = curNode.childs.get(word);
			}
		}
		
		// 결과 출력
		go(root, 0);
		
		bw.flush();
		
		return;
	}
	
	
	/**
	 * DFS의 형태
	 * @param curNode 현재 처리중인 노드
	 * @param depth 굴 깊이
	 * @throws IOException
	 */
	static void go(Node curNode, int depth) throws IOException {
		for (String key : curNode.childs.keySet()) {
			printDepth(depth);
			bw.append(key + '\n');
			go(curNode.childs.get(key), depth + 1);
		}
	}
	
	static void printDepth(int depth) throws IOException {
		for (int i = 0; i < depth; i++) {
			bw.append("--");
		}
	}
	
}