백준 1976 여행 가자 Java [분리집합, Union find]

https://www.acmicpc.net/problem/1976

1976번: 여행 가자

 

 

분리 집합 문제로 여행 계획에 있는 모든 도시들이 하나의 집합에 속해 공통 부모를 가지고 있는지 확인하면 된다

입력에서 1이 나온 부분을 i번째 도시와 j번째 도시를 부모가 같게 만든다

i번째 도시의 i보다 작은 숫자의 도시들은 이전에 중복 체크되니까 j의 시작을 i + 1로 하면 시간을 절반으로 단축할 수 있다

예제의 경우 그림에 빨간 선으로 표현한 것처럼 3번만 확인하면 된다

 

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

public class Main {
    static int[] parents;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int N = Integer.parseInt(br.readLine());
        int M = Integer.parseInt(br.readLine());
        parents = new int[N+1];
        for (int n = 1; n <= N; n++) parents[n] = n;
        for (int i = 0; i < N; i++) {
            int[] cities = getArr(br.readLine());
            for (int j = i + 1; j < N; j++) {
                if (cities[j] == 0) continue;
                union(i + 1, j + 1);
            }
        }
        int[] plan = getArr(br.readLine());
        String result = "YES";
        int parent = find(plan[0]);
        for (int m = 1; m < M; m++) {
            if (parent != find(plan[m])) {
                result = "NO";
                break;
            }
        }
        System.out.println(result);
    }

    static int[] getArr(String str) {
        return Arrays.stream(str.split(" ")).mapToInt(Integer::parseInt).toArray();
    }

    static int find(int x) {
        if (parents[x] == x) return x;
        return parents[x] = find(parents[x]);
    }

    static void union(int x, int y) {
        int px = find(x);
        int py = find(y);
        if (px != py) parents[px] = py;
    }
}