[Java] SWEA 1248 공통조상

SW Expert Academy

문제풀이

키포인트는 이진 트리를 이해하고 응용할 수 있는지와 트리의 성질을 이용해 답을 도출해나갈 수 있는지 두 가지로 보았다. 트리의 성질을 더욱 잘 이해하고 있어야 하는 이유가 있는 문제였다.

 

트리를 활용한 기본 문제들은 입력에서 자식 노드의 인덱스들을 parentIdx * 2 (왼쪽 자식 인덱스), parentIdx * 2 + 1 (오른쪽 자식 인덱스) 그대로를 주는 문제들이 대부분이었다. 다시 말하자면 부모 인덱스가 무조건 자식 인덱스보다 작았었다. 하지만 이 문제의 입력값을 보자면 부모 인덱스가 자식 인덱스보다 값이 큰 경우도 확인해볼 수 있다.

 

그래서 노드 클래스를 이용하여 노드 객체마다 부모인덱스, 왼쪽과 오른쪽 자식인덱스를 가질 수 있도록 했다.

 

또한 문제에서 요구하고 있는 것이 두 가지이다. 첫째는 입력으로 주어지는 Node1, Node2의 공통 부모 인덱스와 둘째는 그것을 루트 노드로 하는 서브 트리의 크기를 구하는 것이다. 처음엔 두번째 요구를 Node1, Node2의 공통 부모로부터 시작하여 Node1, Node2 까지의 크기를 구하는 것이라고 문제를 잘못 읽어서 구상만하다 이상함을 느껴 문제를 다시 읽고 파악했다.

 

요구 기능

1. 공통 부모 찾기

두 노드의 공통 부모를 찾는데 공통 부모 중에서도 Level이 가까운 공통 부모를 찾아야한다. 두 노드로부터 부모 노드를 찾으면서 방문 체크를 한다. 만약 다른 한 노드에서 부모노드를 찾아올라가는데 이미 방문이 되어있다면 두 노드의 공통 부모인 것이다. 또한 문제의 입력에서 전체 트리에서 루트는 항상 1번이라고 되어있다. 이것을 이용하면 된다.

2. 공통 부모를 루트로 하는 서브트리의 크기 구하기

공통 부모를 루트로 두고 모든 노드를 세야 한다. 전위 순회를 하면 된다.

 

코드 및 설명

1. 노드 클래스

이 문제는 노드에 담긴 값과는 관련이 없으므로 따로 data를 담는 변수는 주지 않았다. 또한 인덱스들도 입력에서 주어지면서 생성되기 때문에 기본 생성자로 객체를 만들도록 했다.

static class Node {
    int parentIdx, leftChildIdx, rightChildIdx;
}

 

2. 입력

전역변수에는 입력받을 값들과 노드 클래스 배열인 tree, boolean 타입의 방문 배열이 있다. 이 방문 배열은 입력 받은 두 노드들의 공통 부모를 추적해 나가는데 사용될 것이다. 입력값에 맞춰서 트리를 구성하는데 있어서 중요한 것이 있다. 이진 트리이기 때문에 부모 노드는 최대 두개의 자식을 갖는다. 그래서 편의상 왼쪽부터 채우는 것으로 구현하였는데 사실 방향은 상관 없고 단지 어느 한쪽에 자식노드가 있다면 다른 방향에 자식 노드를 가질 수 있도록 하는 것이 중요하다.

public class Solution {

    static int V, E, Node1, Node2, size;
    static Node[] tree;
    static boolean[] visited;
    
    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;
        
        // 테스트 케이스 입력
        int T = Integer.parseInt(br.readLine());
        for (int t = 1; t <= T; t++) {
            st = new StringTokenizer(br.readLine());
            
            V = Integer.parseInt(br.readLine());
            E = Integer.parseInt(br.readLine());
            Node1 = Integer.parseInt(br.readLine());
            Node2 = Integer.parseInt(br.readLine());
            
            tree = new Node[V + 1];
            visited = new boolean[V + 1];
            
            // 트리 초기화
            for (int i = 1; i <= V; i++) {
                tree[i] = new Node();
            }
            
            st = new StringTokenizer(br.readLine());
            for (int i = 0; i < E; i++) {
                int parent = Integer.parseInt(st.nextToken());
                int child = Integer.parseInt(st.nextToken());
                
                // 이진 트리이므로 부모의 왼쪽 자식이 없으면 왼쪽부터 채우고
                // 있다면 오른쪽에 채우는 식으로 트리를 구성한다.
                // 또한 부모 노드를 추적해야 하므로 child 노드엔 부모 인덱스를 저장해준다.
                if (tree[parent].leftChildIdx == 0) tree[parent].leftChildIdx = child;
                else tree[parent].rightChildIdx = child;
                
                tree[child].parentIdx = parent;
            }
        }
    }
}

 

 

3. 공통 부모 찾기

입력에서 항상 루트 정점은 1이라고 명시했다. Node1과 Node2가 부모노드를 찾아 끝까지 가게 되면 결국 공통 부모는 1인 것이다. 그렇기 때문에 공통 부모를 1로 두고 두 노드가 1이 아닐때 까지 부모노드로 값을 갱신해 나가고 방문한 부모노드를 찾게 되면 그 노드를 공통부모로 두도록 하면 된다. 코드를 보면 바로 이해가 될 것이다.

 

int commonParent = 1;

while (true) {
    if (Node1 != 1) {
        int parent = tree[Node1].parentIdx;

        // 방문한 부모 노드라면 공통 부모이므로 종료한다.
        if (visited[parent]) {
            commonParent = parent;
            break;
        }
        
        // 아니라면 방문 체크를 하고 새로 찾은 부모노드를 a로 저장하여 반복문을 수행한다.
        visited[parent] = true;
        Node1 = parent;
    }

    if (Node2 != 1) {
        int parent = tree[Node2].parentIdx;

        if (visited[parent]) {
            commonParent = parent;
            break;
        }

        visited[parent] = true;
        Node2 = parent;
    }
}

 

4. 공통 부모를 루트 노드로 하는 서브 트리의 크기 구하기

전위, 중위, 후위 순회를 알고 있다면 제일 간단한 기능이다. 공통 부모의 인덱스를 넣어 전위 순회를 돌리는데 재귀의 실행부에서 size를 1씩 증가시켜주기만 하면 된다.

 

static void preOrder(int idx) {
    size++;
    if (tree[idx].leftChildIdx != 0) preOrder(tree[idx].leftChildIdx);
    if (tree[idx].rightChildIdx != 0) preOrder(tree[idx].rightChildIdx);
}

---

전체 코드

 

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

public class Solution {

    static int V, E, Node1, Node2, size;
    static Node[] tree;
    static boolean[] visited;

    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;

        int T = Integer.parseInt(br.readLine());
        for (int t = 1; t <= T; t++) {
            st = new StringTokenizer(br.readLine());

            V = Integer.parseInt(st.nextToken());
            E = Integer.parseInt(st.nextToken());
            Node1 = Integer.parseInt(st.nextToken());
            Node2 = Integer.parseInt(st.nextToken());

            tree = new Node[V + 1];
            visited = new boolean[V + 1];

            for (int i = 1; i <= V; i++) {
                tree[i] = new Node();
            }

            st = new StringTokenizer(br.readLine());
            for (int i = 0; i < E; i++) {
                int parent = Integer.parseInt(st.nextToken());
                int child = Integer.parseInt(st.nextToken());

                if (tree[parent].leftChildIdx == 0) tree[parent].leftChildIdx = child;
                else tree[parent].rightChildIdx = child;

                tree[child].parentIdx = parent;
            }

            int commonParent = 1;

            while (true) {
                if (Node1 != 1) {
                    int parent = tree[Node1].parentIdx;

                    if (visited[parent]) {
                        commonParent = parent;
                        break;
                    }

                    visited[parent] = true;
                    Node1 = parent;
                }

                if (Node2 != 1) {
                    int parent = tree[Node2].parentIdx;

                    if (visited[parent]) {
                        commonParent = parent;
                        break;
                    }

                    visited[parent] = true;
                    Node2 = parent;
                }
            }

            size = 0;
            preOrder(commonParent);

            System.out.println("#" + t + " " + commonParent + " " + size);
        }
    }

    static void preOrder(int idx) {
        size++;
        if (tree[idx].leftChildIdx != 0) preOrder(tree[idx].leftChildIdx);
        if (tree[idx].rightChildIdx != 0) preOrder(tree[idx].rightChildIdx);
    }

    static class Node {
        int parentIdx, leftChildIdx, rightChildIdx;
    }
}

---

정리

트리 응용 문제를 처음 제대로 풀어봐서 기억하기 위해 정리를 해두었다. 자료구조가 달라지더라도 성질을 이해한다면 얼마든지 응용이 가능하다는 것을 몸소 깨닫게 됐다.