[백준] 1976 여행 가자 - 자바

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

1976번: 여행 가자

문제 풀이

union-find 알고리즘을 적용하여 마지막의 주어진 도시들이 같은 그룹에 속해있는지를 비교하면 쉽게 해결할 수 있는 문제다.

 

union-find

private static void union(int a, int b) {
    int rootA = find(a);
    int rootB = find(b);
    if (rootA != rootB) {
        parent[rootA] = rootB;
    }
}

private static int find(int n) {
    if (parent[n] != n)
        parent[n] = find(parent[n]);
    return parent[n];
}

 

이 문제에서 가장 중요한 부분은 다음 아래 코드이다.

boolean canTravel = true;
        for (int i = 1; i < M; i++) {
//            if (parent[travelCities[i]] != parent[travelCities[i - 1]]) {
//                canTravel = false;
//                break;
//            }
            if (find(travelCities[i]) != find(travelCities[i - 1])) {
                canTravel = false;
                break;
            }
        }
        System.out.println(canTravel ? "YES" : "NO");

같은 그룹의 속한 도시인지 확인하기 위해 직접 parent[..]를 참고하면 원하는 결과가 안 나오는 경우가 생긴다.

따라서 find()함수를 통해 아직 루트 노드가 업데이트 시켜줘야 한다. 해당 내용은 이 글에 예시를 통해 설명이 되어있다.

 

코드

package 백준;

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

public class b1979 {
    static int[] parent;

    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());

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

        for (int from = 1; from <= N; from++) {
            StringTokenizer st = new StringTokenizer(br.readLine());
            for (int to = 1; to <= N; to++) {
                int isConnected = Integer.parseInt(st.nextToken());
                if (isConnected == 1)
                    union(from, to);
            }
        }

        int[] travelCities = Arrays.stream(br.readLine().split(" ")).mapToInt(Integer::parseInt).toArray();

        boolean canTravel = true;
        for (int i = 1; i < M; i++) {
//            if (parent[travelCities[i]] != parent[travelCities[i - 1]]) {
//                canTravel = false;
//                break;
//            }
            if (find(travelCities[i]) != find(travelCities[i - 1])) {
                canTravel = false;
                break;
            }
        }
        System.out.println(canTravel ? "YES" : "NO");
    }

    private static void union(int a, int b) {
        int rootA = find(a);
        int rootB = find(b);
        if (rootA != rootB) {
            parent[rootA] = rootB;
        }
    }

    private static int find(int n) {
        if (parent[n] != n)
            parent[n] = find(parent[n]);
        return parent[n];
    }
}