[BOJ][Java]Q2233 - 사과나무

2233번: 사과나무

문제 설명

트리에서 썩은 사과 정점  X, Y가 주어지고,

정점 하나를 제거해서 X, Y를 한번에 없애려면 어떤 정점을 제거해야 하는지에 대한 문제이다.

 

말은 어렵지만 사과나무는 뒤집힌 트리 모양이고,

X, Y를 한 번에 제거할 수 있는 정점은 가장 가까운 공통조상이다.

가장 가까운 공통조상을 구하는 법은 아래 문제에서 풀이한 적이 있다.

[BOJ][Java]Q3584 - 가장 가까운 공통 조상

 

여기까지는 간단하지만 어려웠던 이유는,

트리의 정보가 'dfs 순서의 이진 수열'로 주어지기 때문이다.

정점을 방문할 때 0, 더이상 타고 내려갈 자식이 없어서 부모로 돌아올 때 1로 표시한다.

 

따라서 위와 같은 트리 모양은 dfs 순회했을 때 0001011011이다.

 

여기서 방문할 때의 수열 인덱스를 i, 리턴될때의 수열 인덱스를 j라고 했을 때,

제거해야 할 사과(정점)의 i, j 를 출력하는 것이 문제에서 요구하는 바이다.

 

---

문제 풀이

주어진 수열을 통해 역으로 트리 모양을 유추해내야 하는데,

스택을 이용해서 구현하였다.

 

우선 Pos클래스를 생성했다.

static class Pos{
    int i,j;
}

i는 방문 위치, j는 리턴 위치이다.

 

그리고 Pos배열 pos를 생성했다.

Pos pos[] = new Pos[N+1];
for(int i = 0; i<N+1;i++){
    pos[i] = new Pos();
}

pos[1].i는 1번 정점이 방문되는 수열 인덱스, 

pos[1].j는 1번 정점이 리턴되는 수열 인덱스이다.

 

정점 번호는 int형 변수 node를 이용해
1부터 하나씩 증가시키며 임의로 부여해줄 것이다.

 

 자, 이제 for문으로 arr배열에 저장해둔 이진수열을 돌며 트리정보를 저장한다.

우리가 이 for문에서 할 것은 3가지이다.

• 각 정점의 방문인덱스 i, 반환인덱스 j 찾기
• parent 배열 채우기(각 정점의 부모 찾기) - 공통조상 찾을 때 필요
• 썩은 사과 2개의 정점 번호 찾기 - 알아야 그 정점의 공통조상을 찾음

 

우선 썩은 사과 2개의 정점 번호를 찾기 위해

'0'이든 '1'이든 현재 인덱스가 미리 받아둔 썩은 사과의 인덱스 x, y와 같다면 저장한다.

 

'0'이면

pos[node]의 i값에 현재 for문 인덱스를 저장하고,

stack에 node를 push하며, node 1 증가시킴 

 

'1'이면

stack에서 pop하여 가장 최근에 방문한 정점을 리턴한다. 

그 값을 now라고 했을 때, pos[now].j에 현재 인덱스를 저장한다.

여기서 한 가지 더 따져주어야 한다.

1일 경우 now에서 리턴되는 순간이므로, 다음에 pop될 정수는 현재 정점의 부모이다.

따라서 한 번 더 pop하여 그 정수를 parent 배열에 저장한 후 다시 스택에 push해준다.

        Stack<Integer> stack = new Stack<>();
        int node = 1;
        for(int index = 1; index<str.length()+1;index++){
            if(arr[index] == '0') {
                //썩은 사과 정점 번호 찾기
                if(x==index)    x = node;
                if(y==index)    y = node;

                pos[node].i = index;
                stack.push(node++);
            }else if(arr[index] == '1'){
                int now = stack.pop();


                //썩은 사과 정점 번호 찾기
                if(x==index)    x = now;
                if(y==index)    y = now;

                pos[now].j = index;
                if(stack.isEmpty()) break;

                //부모찾기
                int nowParent = stack.pop();    //하나 더 pop해서 부모 찾고
                parent[now] = nowParent;
                stack.push(nowParent);  //다시 push해줌
            }
        }

 

이제 트리 정보가 정리되었으니 parent배열을 이용해 공통조상을 찾아주면 된다.

찾기에 앞서, x와 y가 같은 정점이면 그냥 그 정점만 제거하면 되므로 x의 i, j를 출력하고 끝낸다.

        if(x==y){
            System.out.println(pos[x].i+" "+pos[x].j);
            return;
        }

공통조상을 찾는 방법은

x의 부모를 while문으로 타고 올라가며 리스트에 모두 저장한다.(x 자기자신 포함)

그리고 y의 부모를 타고 올라가며 해당 list가 y의 부모를 포함하고 있다면 반복문을 멈추고

해당 위치에서 i, j를 출력하면 된다.

        ArrayList<Integer> xParents = new ArrayList<>();
        while(x!=0){
            xParents.add(x);
            x = parent[x];
        }
        int result = 0;
        while(y!=0){
            if(xParents.contains(y)){
                result = y;
                break;
            }
            y = parent[y];
        }
        System.out.println(pos[result].i+" "+pos[result].j);

---

전체 코드 

package 트리;

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

public class Q2233_사과나무 {
    /*[골드1]Q2233 - 사과나무*/
    static int N;
    public static void main(String[] args) throws IOException {
        BufferedReader br =new BufferedReader(new InputStreamReader(System.in));
        N = Integer.parseInt(br.readLine());

        int parent[] = new int[N+1];
        Pos pos[] = new Pos[N+1];
        for(int i = 0; i<N+1;i++){
            pos[i] = new Pos();
        }
        String str = br.readLine();
        char arr[] = new char[str.length()+1];
        for(int i = 1; i<str.length()+1;i++){
            arr[i] = str.charAt(i-1);
        }

        StringTokenizer st = new StringTokenizer(br.readLine());
        int x = Integer.parseInt(st.nextToken());
        int y = Integer.parseInt(st.nextToken());

        Stack<Integer> stack = new Stack<>();
        int node = 1;
        for(int index = 1; index<str.length()+1;index++){
            if(arr[index] == '0') {
                //썩은 사과 정점 번호 찾기
                if(x==index)    x = node;
                if(y==index)    y = node;

                pos[node].i = index;
                stack.push(node++);
            }else if(arr[index] == '1'){
                int now = stack.pop();


                //썩은 사과 정점 번호 찾기
                if(x==index)    x = now;
                if(y==index)    y = now;

                pos[now].j = index;
                if(stack.isEmpty()) break;

                //부모찾기
                int nowParent = stack.pop();    //하나 더 pop해서 부모 찾고
                parent[now] = nowParent;
                stack.push(nowParent);  //다시 push해줌
            }
        }

        if(x==y){
            System.out.println(pos[x].i+" "+pos[x].j);
            return;
        }

        ArrayList<Integer> xParents = new ArrayList<>();
        while(x!=0){
            xParents.add(x);
            x = parent[x];
        }
        int result = 0;
        while(y!=0){
            if(xParents.contains(y)){
                result = y;
                break;
            }
            y = parent[y];
        }
        System.out.println(pos[result].i+" "+pos[result].j);
    }
    static class Pos{
        int i,j;

        public Pos() {
        }
    }
}