[SWEA | Java] 1248번 공통조상

문제

https://swexpertacademy.com/main/code/problem/problemDetail.do?contestProbId=AV15PTkqAPYCFAYD

SW Expert Academy

티어

소요 시간

35분

분류

트리

코드

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

public class Solution {

    static int v, e;
    static int v1, v2;
    static Node[] nodes;

    static public class Node {
        int element;
        Node parent;
        Node leftChild;
        Node rightChild;

        public Node(int element, Node parent, Node leftChild, Node rightChild) {
            this.element = element;
            this.parent = parent;
            this.leftChild = leftChild;
            this.rightChild = rightChild;
        }
    }

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

        int t = Integer.parseInt(st.nextToken());
        for (int test = 1; test <= t; test++) {
            st = new StringTokenizer(br.readLine(), " ");
            v = Integer.parseInt(st.nextToken());
            e = Integer.parseInt(st.nextToken());
            v1 = Integer.parseInt(st.nextToken());
            v2 = Integer.parseInt(st.nextToken());
            nodes = new Node[v + 1];

            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 (nodes[parent] == null) {
                    nodes[parent] = new Node(parent, null, null, null);
                }
                if (nodes[child] == null) {
                    nodes[child] = new Node(child, nodes[parent], null, null);
                }
                else {
                    nodes[child].parent = nodes[parent];
                }
                
                if (nodes[parent].leftChild == null) {
                    nodes[parent].leftChild = nodes[child];
                }
                else {
                    nodes[parent].rightChild = nodes[child];
                }

            }

            boolean[] visited = new boolean[v + 1];
            getParents(nodes[v1], visited);

            Node nearParent = findNearParent(nodes[v2], visited);
            int cnt = findChildrenCnt(nearParent);

            System.out.println("#" + test + " " + nearParent.element + " " + cnt);
        }
    }

    private static int findChildrenCnt(Node nearParent) {
        int cnt = 1;
        if (nearParent.leftChild != null) {
            cnt += findChildrenCnt(nearParent.leftChild);
        }
        if (nearParent.rightChild != null) {
            cnt += findChildrenCnt(nearParent.rightChild);
        }
        return cnt;
    }

    private static Node findNearParent(Node target, boolean[] visited) {
        Node parent = nodes[1];
        for (Node x = target; x != null; x = x.parent) {
            if (visited[x.element]) {
                parent = x;
                break;
            }
        }
        return parent;
    }

    private static void getParents(Node target, boolean[] visited) {
        for (Node x = target; x != null; x = x.parent) {
            visited[x.element] = true;
        }
    }
}

코드 설명

너무 명확한 트리문제였기 때문에 문제에 최적화된 풀이 방법을 푸는데 주력했다.

문제에서 가장 중요하게 요구하는 것은 가까운 공통조상을 찾는 것이다.

조상을 찾기 가장 쉬운 방법은 모든 요소들을 Node라는 클래스로 두고 parent와 child에 모두 접근할 수 있게 하는 것이다.

그런데 문제의 입력이 순서대로가 아닌 뒤죽박죽 들어온다는 점, 따라서 조회가 잦다는 점에서 배열을 선택하게 되었다.

트리의 크기가 방대하지 않고 조회 위주의 문제이기 때문에 조회에 O(1)의 시간복잡도가 소요되는 배열이 가장 적합한 자료구조라고 생각했다.