[BOJ] 14725. 개미굴 - JAVA

문제 링크

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

14725번: 개미굴

 

풀이

문제 유형 : 트라이

* 트라이의 기본 코드에 대한 설명 : https://github.com/J00HUI/Algorithm-Codes/blob/master/trie.md

 

이 문제에서 기본 트라이 문제와 다르게 신경써야 할 점은

1. Map 내에 알파벳 순으로 정렬

2. 트라이 내의 원소들을 원하는 형태로 출력

 

이다. 

먼저 1번은 HashMap 대신 TreeMap 을 써서 해결할 수 있다.

Map<String, Node> childNode = new TreeMap<>();

추가로 보통은 Key 값에 char 형을 넣지만, 두 번째 예제의 경우 문자열이 하나의 노드에 들어가야 하기 때문에 String 을 사용했다.

 

2번 원하는 형태로 출력에서 막혀서 블로그를 참고했다.

Trie 클래스에 print() 메서드를 추가해서 자식노드를 탐색하는 과정에서 문제의 출력형식대로 출력하도록 했다.

print() 의 매개변수로 현재 노드의 높이 정보(floor) 를 넣어 높이의 수 만큼  '--' 가 출력 되도록 한다.

예를 들어, 높이가 1이면 '--' 만 출력되고, 높이가 2라면 '----'가 출력되어야 한다. 

 

void print() {
    print(rootNode, 0);
}
 
void print(Node thisNode, int floor) {
    Set<String> set = thisNode.childNode.keySet();        // thisNode 의 자식노드(Map)의 Key 들을 set 으로 가져옴
    Iterator<String> iter = set.iterator();                // iterator 선언
     
    while(iter.hasNext()) {
        String str = iter.next();                        // Key 
        
        for(int i=0; i<floor; i++) {                    // 높이(depth) 만큼 '--' 을 출력
            sb.append("--");
        }
        
        sb.append(str).append("\n");                    // 키 출력
        
        Node childNodes = thisNode.childNode.get(str);    // thisNode 의 Map 중에서 현재 키를 찾아 노드로 반환
        print(childNodes, floor+1);                        // 반환된 노드의 자식 노드들 재귀로 탐색, floor 1 증가(floor = 높이(depth))
    }
}

 

코드

import java.io.*;
import java.util.*;
 
public class Main_bj_14725_개미굴 {
    
    static StringBuilder sb = new StringBuilder();
    
    static class Node {
        Map<String, Node> childNode = new TreeMap<>();
        boolean endOfStr;
    }
    
    static class Trie {
        
        Node rootNode = new Node();
        
        void insert(String str) {
            StringTokenizer st = new StringTokenizer(str, " ");
            
            int K = Integer.parseInt(st.nextToken());
            Node root = this.rootNode;
            
            for(int i=0; i<K; i++) {
                root = root.childNode.computeIfAbsent(st.nextToken(), key -> new Node());
            }
            
            root.endOfStr = true;
        }
        
        void print() {
            print(rootNode, 0);
        }
        
        void print(Node thisNode, int floor) {
            Set<String> set = thisNode.childNode.keySet();
            Iterator<String> iter = set.iterator();
             
            while(iter.hasNext()) {
                String str = iter.next();
                
                for(int i=0; i<floor; i++) {
                    sb.append("--");
                }
                
                sb.append(str).append("\n");
                
                Node childNodes = thisNode.childNode.get(str);
                print(childNodes, floor+1);
            }
        }
        
    }
 
    public static void main(String[] args) throws Exception {
//        System.setIn(new FileInputStream("res/input_bj_14725.txt"));
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        
        int N = Integer.parseInt(br.readLine());
        StringTokenizer st;
        
        Trie trie = new Trie();
        for(int i=0; i<N; i++) {
            trie.insert(br.readLine());     // 2 KIWI BANANA
        }
        
        trie.print();
        System.out.println(sb.toString());
        br.close();
    }
 
}