[BOJ] 9489 : 사촌(Java)

문제

증가하는 정수 수열을 이용해서 트리를 만드는 방법은 다음과 같다.

• 첫 번째 정수는 트리의 루트 노드이다.
• 다음에 등장하는 연속된 수의 집합은 루트의 자식을 나타낸다. 이 집합에 포함되는 수의 첫 번째 수는 항상 루트 노드+1보다 크다.
• 그 다음부터는 모든 연속된 수의 집합은 아직 자식이 없는 노드의 자식이 된다. 그러한 노드가 여러 가지 인 경우에는 가장 작은 수를 가지는 노드의 자식이 된다.
• 집합은 수가 연속하지 않는 곳에서 구분된다.

예를 들어, 수열 1 3 4 5 8 9 15 30 31 32를 위의 규칙을 이용해 트리를 만들면 아래 그림과 같이 된다.

두 노드의 부모는 다르지만, 두 부모가 형제(sibling)일 때 두 노드를 사촌이라고 한다.

수열 특정 노드 번호 k가 주어졌을 때, k의 사촌의 수를 구하는 프로그램을 작성하시오.

입력

입력은 여러 개의 테스트 케이스로 이루어져 있다. 각 테스트 케이스의 첫째 줄에는 노드의 수 n과 사촌의 수를 구해야 하는 노드의 번호 k가 주어진다. (1 ≤ n ≤ 1,000, 1 ≤ k ≤ 1,000,000) 다음 줄에는 총 n개의 수가 주어지며, 모든 수는 1보다 크거나 같고, 1,000,000보다 작거나 같다. 입력으로 주어지는 수열은 항상 증가한다. k는 항상 수열에 포함되는 수이다.
입력의 마지막 줄에는 0이 두 개 주어진다.

출력

각 테스트 케이스 마다, k의 사촌의 수를 출력한다.

문제 링크 : https://www.acmicpc.net/problem/9489

---

💡 풀이

주어진 조건에 만족하도록 트리를 만드는 것이 제일 중요한 풀이 과정이다. 

우선, 두 번째 줄에 들어오는 노드들을 연속된 형태인지 비교하며, 노드들을 알맞게 연결해준다. 

주의할 점은 트리를 구현하기 위해 ArrayList를 사용하여 인접 리스트를 이용하면 메모리 초과가 발생한다. 따라서, 각 인덱스(노드)의 부모를 저장하는 parent 배열과 각 인덱스(노드)의 자식의 개수를 저장하는 child 배열, 두 가지를 선언하여 풀었다. 

트리를 구현해야 한다고 해서 반드시 인접 리스트를 사용해야 할 필요는 없다는 것을 깨달았다..!

 

전체적인 코드는 다음과 같다. 

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.StringTokenizer;

public class Q9489 {
    static int N, K;
    static ArrayList<Integer> node;
    
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;

        while (true) {
            st = new StringTokenizer(br.readLine());
            N = Integer.parseInt(st.nextToken());
            K = Integer.parseInt(st.nextToken());

            if (N == 0 && K == 0) {
                break;
            }

            st = new StringTokenizer(br.readLine());
            node = new ArrayList<>();
            for (int i = 0; i < N; i++) {
                node.add(Integer.parseInt(st.nextToken()));
            }

            // 트리 만들기
            int cur_root = 0; // 현재 연결하고자 하는 부모 노드
            int[] parent = new int[N + 1]; // 각 인덱스의 부모 노드
            int[] child = new int[N + 1]; // 인덱스의 자식의 개수
            parent[cur_root] = -1;
            for (int i = 1; i < N; i++) {
                int cur_node = node.get(i);
                int prev_node = node.get(i - 1);

                // 연속되지 않은 수열이 오면 새로운 부모 노드를 찾아 연결
                if (cur_node - prev_node != 1) {
                    cur_root = findNext(child);
                }
                parent[i] = cur_root;
                child[cur_root]++;
            }

            int cousin = findCousin(parent, child);
            System.out.println(cousin);

        }

    }

    static int findNext(int[] child) {
        // 자식이 없는 최소 노드 찾아 반환하기
        for (int i = 0; i <= N; i++) {
            if (child[i] == 0) {
                return i;
            }
        }
        return -1;
    }

    static int findCousin(int[] parent, int[] child) {
        int index = node.indexOf(K);
        int K_parent = parent[index]; // K의 부모 와는 달라야하고,
        if (K_parent == -1) {
            return 0;
        }
        int K_parent_root = parent[K_parent]; // K의 부모의 부모와는 같아야 한다.
        if (K_parent_root == -1) {
            return 0;
        }

        int ans = 0;
        ArrayList<Integer> list = new ArrayList<>();
        for (int i = 0; i < N; i++) {
            if (parent[i] == K_parent_root) {
                list.add(i);
            }
        }

        for (int i = 0; i < list.size(); i++) {
            int cur = list.get(i);
            if (cur != K_parent) {
                ans += child[cur];
            }
        }

        return ans;
    }
}