[백준] 9489번 - 사촌 Python, 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의 사촌의 수를 출력한다.

예제 입력 1

---

10 15
1 3 4 5 8 9 15 30 31 32
12 9
3 5 6 8 9 10 13 15 16 22 23 25
10 4
1 3 4 5 8 9 15 30 31 32
0 0

예제 출력 1

---

5
1
0

풀이

---

Python

import sys
input = sys.stdin.readline

while True:
    n,k = map(int,input().rstrip().split())
    if n==0 and k == 0:
        break
    a = [-1] + list(map(int,input().rstrip().split()))
    parent = [0]*(n+1)
    parent[0] = -1
    target = 0
    idx = -1
    for i in range(1,n+1):
        if a[i] == k:
            target = i
        if a[i] != a[i-1]+1:
            idx+=1
        parent[i] = idx
    answer = 0
    for i in range(1,n+1):
        if parent[i] != parent[target] and parent[parent[i]] == parent[parent[target]]:
            answer+=1
    print(answer)

Java

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

public class Main {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringBuilder sb = new StringBuilder();
        while(true){
            StringTokenizer st = new StringTokenizer(br.readLine());
            int n = Integer.parseInt(st.nextToken());
            int k = Integer.parseInt(st.nextToken());
            if(n==0 && k==0) break;
            int target = 0;
            int[] arr = new int[n+1];
            int[] parent = new int[n+1];
            int idx = -1;
            arr[0] = -1;
            parent[0] = -1;
            st = new StringTokenizer(br.readLine());
            for(int i = 1; i <=n; i++) {
                arr[i] = Integer.parseInt(st.nextToken());
                if(arr[i] == k) target = i;
                if(arr[i] != arr[i-1]+1) idx++;
                parent[i] = idx;
            }
            int answer = 0;
            for(int i = 1; i<=n; i++) {
                if(parent[i] != parent[target] && parent[parent[i]] == parent[parent[target]]) answer++;
            }
            sb.append(answer+"\n");
        }
        System.out.println(sb.toString());
    }
}

정리

---

• 노드 번호가 2이상 차이나는 부분을 집합으로 구분하여 부모노드를 구해준다.
  1,3,4,5,8,9,15,30,31,32 이렇게 나눠져서 각 노드들의 부모노드들을 구해주면 된다.
• 그 다음 target의 부모와 다르면서 target의 부모의 부모가 일치하면 카운트 해준다.