[백준 4256] 트리(Java)

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

이 문제는 트리의 성질을 이용하여 풀 수 있습니다.

우선 전위 순회의 경우 루트 노드부터 시작하여 리프 노드까지 탐색합니다. 따라서 가장 먼저 루트 노드가 나오고, 그 다음 인덱스가 루트 노드의 자식 노드 순으로 내려갑니다.

반면 중위 순회의 경우 왼쪽에서 오른쪽으로 탐색합니다. 따라서 루트 노드값을 안다면 그 값을 기준으로 왼쪽은 왼쪽 자식 노드, 오른쪽은 오른쪽 자식 노드가 됩니다.

이를 동시에 적용하면 다음과 같이 트리를 구현할 수 있습니다.

1. 전위 순회를 진행하면서 현재의 부모 노드를 체크
2. 부모 노드를 기준으로 중위 순회 배열을 두 개로 나눔(왼쪽, 오른쪽)
3. 각각의 나눈 배열을 재귀함수를 이용하여 탐색한 결과를 트리의 자식 노드로 설정

이를 조금 더 구체적으로 설명해보겠습니다.

우선 가장 먼저 전위 순회 배열의 0번 인덱스를 탐색합니다. 이 값은 트리 전체의 루트 노드가 됩니다. 이 루트 노드와 같은 값을 가지는 중위 순회 배열의 인덱스를 구합니다. 이 인덱스를 구하는 이유는 이 값을 기준으로 중위 순회 배열의 왼쪽은 왼쪽 자식 노드, 오른쪽은 오른쪽 자식 노드이기 때문입니다.

위에서 구한 정보를 바탕으로 트리를 구현합니다. 재귀함수는 트리의 노드를 반환하여 부모 노드의 값을 가지는 노드를 생성한 후 왼쪽 자식 노드를 위에서 구한 왼쪽 자식 노드, 오른쪽 자식 노드를 오른쪽 자식 노드값으로 치환합니다. 이 때 재귀호출을 이용하여 위의 과정을 반복합니다.

이 때 전위 순회의 어떤 인덱스부터 탐색할지를 정해야 합니다. 중위 순회에서 왼쪽의 원소들의 개수가 곧 왼쪽 자식 노드의 개수이기 때문에 전위 순회에서 첫 번째 노드를 제외하고 왼쪽의 원소들의 개수만큼 건너뛰면 오른쪽 자식 노드가 시작합니다. 이 때 전위 순회 특성상 가장 먼저 오는 인덱스가 오른쪽 자식 노드의 부모 노드이기 때문에 그 인덱스를 가리키게 리턴합니다.

다음은 코드입니다.

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

class Main{
    static class Node{
        int val;
        Node left, right;

        Node(int val){
            this.val = val;
        }
    }

    static int N;
    static StringBuilder sb = new StringBuilder();

    public static void main(String[] args) throws Exception{
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int T = Integer.parseInt(br.readLine());
        while(T-- >0){
            N = Integer.parseInt(br.readLine());
            int[] preorder = new int[N];
            StringTokenizer st = new StringTokenizer(br.readLine());
            for(int i=0;i<N;i++) preorder[i] = Integer.parseInt(st.nextToken());

            st = new StringTokenizer(br.readLine());
            int[] inorder = new int[N];
            for(int i=0;i<N;i++) inorder[i] = Integer.parseInt(st.nextToken());

            Node tree = buildTree(0, preorder, inorder);
            postOrder(tree);
            sb.append("\n");
        }
        System.out.println(sb);
    }

    static void postOrder(Node node){
        if(node == null) return;

        postOrder(node.left);
        postOrder(node.right);
        sb.append(node.val).append(" ");
    }

    static Node buildTree(int idx, int[] preorder, int[] inorder){
        if(idx >= N) return null;

        int root = getIndex(preorder[idx], inorder);
        if(root == -1) return null;
        int[] leftArr = Arrays.copyOfRange(inorder, 0, root);
        int[] rightArr = Arrays.copyOfRange(inorder, root+1, inorder.length);

        Node node = new Node(inorder[root]);
        node.left = buildTree(idx+1, preorder, leftArr);
        node.right = buildTree(idx+root+1, preorder, rightArr);
        return node;
    }

    static int getIndex(int target, int[] inorder){
        for(int i=0;i<inorder.length;i++){
            if(inorder[i] == target) return i;
        }
        return -1;
    }
}