[백준] 1976번 - 여행 가자 (Gold 4)

문제 링크

[백준] 1976번 - 여행 가자 (Gold 4)

문제 설명

전체 도시의 수 $N(\leq 200)$과 여행 계획에 속한 도시의 수 $M(\leq 1000)$이 주어진다.
도시 i와 j가 연결되어있는지 아닌지의 여부와 여행계획도 입력으로 주어진다.
이 때, 여행 계획이 가능한 계획인지 알아내는 프로그램을 작성하시오.

정답 코드 및 설명

역시 유니온 파인드 알고리즘을 사용하는 문제이다.
두 도시가 연결되어 있다면 같은 집합에 넣는다.
여행 계획에 포함된 모든 도시가 같은 집합에 포함되어 있다면, 가능한 여행계획이다.
백준 1717번 - 집합의 표현 (Gold 4)의 소스 코드를 상당 부분 재활용하였다.

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

public class BOJ1976 {
    int n, m;
    int[][] adj;
    int[] path;
    Node[] set;

    class Node {
        int n, level;
        Node parent;

        Node(int n) {
            this.n = n;
            this.level = 0;
            setParent(this);
        }

        void setParent(Node parent) {
            this.parent = parent;
        }
    }

    void input() throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        n = Integer.parseInt(br.readLine());
        m = Integer.parseInt(br.readLine());
        set = new Node[n];
        path = new int[m];
        StringTokenizer st;
        for (int i = 0; i < n; i++) {
            st = new StringTokenizer(br.readLine());
            set[i] = new Node(i);
            for (int j = 0; j < i; j++) {
                int adj = Integer.parseInt(st.nextToken());
                if (adj == 0) continue;
                if (findSet(i) == findSet(j)) continue;
                union(i, j);
            }
        }
        st = new StringTokenizer(br.readLine());
        for (int i = 0; i < m; i++)
            path[i] = Integer.parseInt(st.nextToken()) - 1;
    }

    int findSet(int a) {
        Node curr = set[a];
        while (curr != curr.parent) {
            curr = curr.parent;
        }
        return curr.n;
    }

    String isAble(int[] path) {
        int set = findSet(path[0]);
        for (int i = 1; i < path.length; i++) {
            if (set != findSet(path[i])) return "NO";
        }
        return "YES";
    }

    void union(int a, int b) {
        int aSet = findSet(a);
        int bSet = findSet(b);
        if (aSet == bSet) return;
        int aSetLevel = set[aSet].level;
        int bSetLevel = set[bSet].level;
        if (aSetLevel >= bSetLevel) {
            set[bSet].parent = set[aSet];
            if (aSetLevel == bSetLevel) set[aSet].level++;
            return;
        }
        set[aSet].parent = set[bSet];
    }

    void solution() throws IOException {
        input();
        System.out.println(isAble(path));
    }

    public static void main(String[] args) throws IOException {
        new BOJ1976().solution();
    }
}