[Java] 백준 9934번 : 완전 이진 트리

[Silver I] 완전 이진 트리 - 9934

문제 링크

성능 요약

메모리: 15040 KB, 시간: 180 ms

분류

재귀, 트리

문제 설명

상근이는 슬로베니아의 도시 Donji Andrijevci를 여행하고 있다. 이 도시의 도로는 깊이가 K인 완전 이진 트리를 이루고 있다. 깊이가 K인 완전 이진 트리는 총 2^K-1개의 노드로 이루어져 있다. (아래 그림) 각 노드에는 그 곳에 위치한 빌딩의 번호가 붙여져 있다. 또, 가장 마지막 레벨을 제외한 모든 집은 왼쪽 자식과 오른쪽 자식을 갖는다.

깊이가 2와 3인 완전 이진 트리

상근이는 도시에 있는 모든 빌딩에 들어갔고, 들어간 순서대로 번호를 종이에 적어 놓았다. 한국으로 돌아온 상근이는 도시가 어떻게 생겼는지 그림을 그려보려고 하였으나, 정확하게 기억이 나지 않아 실패했다. 하지만, 어떤 순서로 도시를 방문했는지 기억해냈다.

1. 가장 처음에 상근이는 트리의 루트에 있는 빌딩 앞에 서있다.
2. 현재 빌딩의 왼쪽 자식에 있는 빌딩에 아직 들어가지 않았다면, 왼쪽 자식으로 이동한다.
3. 현재 있는 노드가 왼쪽 자식을 가지고 있지 않거나 왼쪽 자식에 있는 빌딩을 이미 들어갔다면, 현재 노드에 있는 빌딩을 들어가고 종이에 번호를 적는다.
4. 현재 빌딩을 이미 들어갔다 온 상태이고, 오른쪽 자식을 가지고 있는 경우에는 오른쪽 자식으로 이동한다.
5. 현재 빌딩과 왼쪽, 오른쪽 자식에 있는 빌딩을 모두 방문했다면, 부모 노드로 이동한다.

왼쪽 그림에 나와있는 마을이라면, 상근이는 2-1-3 순서대로 빌딩을 들어갔을 것이고, 오른쪽 그림의 경우에는 1-6-4-3-5-2-7 순서로 들어갔을 것이다. 상근이가 종이에 적은 순서가 모두 주어졌을 때, 각 레벨에 있는 빌딩의 번호를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 K (1 ≤ K ≤ 10)가 주어진다.

둘째 줄에는 상근이가 방문한 빌딩의 번호가 들어간 순서대로 주어진다. 모든 빌딩의 번호는 중복되지 않으며, 구간 [1,2^K)에 포함된다.

출력

총 K개의 줄에 걸쳐서 정답을 출력한다. i번째 줄에는 레벨이 i인 빌딩의 번호를 출력한다. 출력은 왼쪽에서부터 오른쪽 순서대로 출력한다.

풀이

소스 코드

---

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

public class Main {
    static ArrayList<ArrayList<Integer>> tree;
    static int[] arr;
    static int k;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        k = Integer.parseInt(br.readLine());
        int nodes = (int) Math.pow(2, k) - 1;
        arr = new int[nodes];

        StringTokenizer st = new StringTokenizer(br.readLine());
        for (int i = 0; i < nodes; i++) {
            arr[i] = Integer.parseInt(st.nextToken());
        }

        tree = new ArrayList<ArrayList<Integer>>();
        for (int i = 0; i < nodes; i++) {
            tree.add(new ArrayList<Integer>());
        }

        if (k == 1) {
            System.out.println(arr[0]);
            return;
        }
        else {
            getAns((nodes-1)/2, 1);
        }

        for (int i = 1; i < k+1; i++) {
            for (int j = 0; j < Math.pow(2, i-1); j++) {
                System.out.print(tree.get(i).get(j) + " ");
            }
            System.out.println();
        }
    }
    static void getAns(int midIdx, int depth) {
        if (depth == k) {
            tree.get(depth).add(arr[midIdx]);
            return;
        }
        tree.get(depth).add(arr[midIdx]);
        getAns(midIdx-((int)Math.pow(2, (k-(depth+1)))), depth+1);
        getAns(midIdx+((int)Math.pow(2, (k-(depth+1)))), depth+1);
    }
}

설명

---

1. Root 노드는 항상 입력의 정 가운데에 위치

   완전 이진 트리(포화 이진 트리)만 주어지는 상황이기 때문에, 입력으로 주어진 값들 정 가운데 있는 값은 항상 Root 노드 값으로 고정임을 알 수 있습니다.

2. 재귀적으로 접근, 각 서브 트리의 루트 역시 가운데에 존재함

   가운데 Root를 기준으로, 왼쪽 서브 트리와 오른쪽 서브 트리 역시 완전 이진 트리입니다. 즉, 1번의 상황을 재귀적으로 적용할 수 있습니다.

3. 재귀적 반복이 끝나는 시점은 트리의 깊이(depth)에 다다랐을 시점

   특정 노드에 대해 왼쪽, 오른쪽 서브 트리에 관한 함수를 재귀적으로 호출하며 depth 값을 늘려가다가, 입력에 주어진 depth와 일치할 때 재귀 반복을 종료합니다.

---

1. 필요한 변수 선언

static ArrayList<ArrayList<Integer>> tree;
static int[] arr; 
static int k; // depth

tree는 ArrayList 로 구성된 2차원 배열입니다. 각 depth마다 노드를 저장하여 정답을 출력하는 용도입니다. arr는 일반 정수 배열입니다. 입력된 정수 노드 값을 저장합니다.

2. 입력 받기, 배열 설정

BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
k = Integer.parseInt(br.readLine());
int nodes = (int) Math.pow(2, k) - 1;
arr = new int[nodes];

StringTokenizer st = new StringTokenizer(br.readLine());
for (int i = 0; i < nodes; i++) {
    arr[i] = Integer.parseInt(st.nextToken());
}

tree = new ArrayList<ArrayList<Integer>>();
for (int i = 0; i < nodes; i++) {
    tree.add(new ArrayList<Integer>());
}

BufferedReader 클래스를 통해 입력을 받습니다. 다만 이번 문제는 입력에서의 시간 딜레이가 유의미한 영향을 주지 않으므로 Scanner 클래스를 사용하여도 큰 문제는 없을 것 같습니다.

k 값을 입력 받고, 편의상 트리의 전체 노드 수를 계산하여 nodes 변수에 저장했습니다.

문제에서 주어진 대로, 깊이가 $k$ 인 완전 이진 트리의 노드 개수는 $2^k - 1$ 입니다.

공백으로 구분된 정수 입력을 받기 위해 StringTokenizer 클래스를 사용합니다.

for 반복을 통해 그냥 정수 배열에 집어 넣습니다.

tree 배열을 초기화합니다. 코드에서는 nodes 를 크기로 하여 초기화하였는데, k 를 크기로 하는 것이 맞습니다. 위 그림처럼 트리를 depth 기준으로 저장하기 때문입니다. 괜히 메모리를 더 잡아먹지 않았나 싶네요..

3. 재귀 함수 구현

static void getAns(int midIdx, int depth) {
        if (depth == k) {
            tree.get(depth).add(arr[midIdx]);
            return;
        }
        tree.get(depth).add(arr[midIdx]);
        getAns(midIdx-((int)Math.pow(2, (k-(depth+1)))), depth+1);
        getAns(midIdx+((int)Math.pow(2, (k-(depth+1)))), depth+1);
    }

기본적으로, 인자로 두 가지를 받습니다. midIdx와 depth.

• midIdx : 함수가 탐색 중인 트리의 Root 노드 인덱스를 가리킵니다.
• depth : 함수가 탐색 중인 현재 뎁스를 나타냅니다.

그 다음으로 탈출 조건을 설정합니다. depth 가 k 와 같은 경우, midIdx를 2차원 배열 tree에 추가합니다. 예시처럼 뎁스가 3인 트리일 때, midIdx 값은 1, 4, 5, 7 일 것입니다. 마지막 뎁스이기 때문에 하위 노드가 없을 뿐이지, 만일 뎁스가 4인 트리였다면 위 값들도 하위 노드를 두 개씩 가지는 서브트리의 Root입니다.

현재 탐색 중인 depth가 k보다 작으면, 다음 작업을 수행합니다.

1. 우선 인자로 들어온 midIdx는 해당 depth의 노드입니다. tree에 추가합니다.
2. 재귀 호출을 시행합니다. 좌측 서브트리, 우측 서브트리에 대해 각각 실행됩니다.

재귀 호출 시 전달되는 midIdx 계산식이 조금 복잡해 보이나, 뜯어보면 쉽습니다. 깊이가 4인 트리를 예로 들어 봅시다.

이해하기 편하도록 각 value를 해당 노드의 depth로 모두 세팅한 트리를 가정합니다.

1. depth = 1 에서 2로 가면서, 초기 midIdx 값은 4만큼 줄어듭니다.
2. 2에서 3으로 가면서, 두번째 값은 2만큼 줄어듭니다.
3. 3에서 4로 가면서, 세번째 값은 1만큼 줄어듭니다.
4. depth = 4일때는 재귀 탈출 조건을 통해 반복이 종료됩니다.

즉, 식을 세우자면 $ \text{midIdx} \pm 2^{k-(\text{depth}+1)} $ 로 볼 수 있겠습니다. 부호에 따라 마이너스면 좌측, 플러스면 우측 서브트리의 루트로 이동하게 됩니다.

depth 인자로는 당연히 다음 뎁스를 탐색하는 것이므로, depth+1 을 전달합니다.

4. 정답 출력

if (k == 1) {
    System.out.println(arr[0]);
    return;
}
else {
    getAns((nodes-1)/2, 1);
}

for (int i = 1; i < k+1; i++) {
    for (int j = 0; j < Math.pow(2, i-1); j++) {
        System.out.print(tree.get(i).get(j) + " ");
    }
    System.out.println();
}

구현한 함수를 main에서 호출합니다.
k가 1일 경우의 예외 처리를 따로 마련해 놓긴 했지만, 구현된 함수의 특성상 1인 경우에도 정상적으로 동작할 것으로 보입니다.

가장 처음의 midIdx 값은 $ \frac{\text{노드 수}-1}{2} $ 로 계산됩니다.

이중 for문을 통해 완성된 tree의 요소들을 전부 출력해주면 끝입니다.