[JAVA] 백준 1976 유니온 파인드

 

 

 

1976번: 여행 가자

문제

도시의 개수가 주어지고 각 도시들이 연결되어 있는지 여부가 주여졌을 때 가려고 하는 여행지를 모두 갈 수 있는지 여부를 출력하시오.

풀이

📌 유니온 파인드

    static void union(int a, int b){
        int x = find(a);
        int y = find(b);

        A[y] = x;
    }
    static int find(int a){
        if(A[a] == a){
            return a;
        }else{
            A[a] = find(A[a]);
            return A[a];
        }
    }

union : 주어진 수들의 집합을 찾아서 y를 x집합에 속하게 한다.

find : 주어진 수의 집합의 가장 상단을 재귀적으로 돌면서 찾아서 담아둔다.

📌 유니온 연산 (주어진 도시들이 연결된 여부)

        for(int i=1; i<N; i++){
            st = new StringTokenizer(br.readLine(), " ");
            for(int j=1; j<N; j++){
                boolean c = st.nextToken().equals("1") ? true : false;
                if(c) union(i, j);
            }
        }

입력으로 1이 들어온 수와 현재 행의 도시가 연결되어 있다는 의미이므로 유니온 연산을 통해 같은 집합으로 둔다.

📌 파인드 연산 (여행 가려는 도시들이 모두 같은 집합에 속하는지 여부를 검사)

        st = new StringTokenizer(br.readLine());
        int target = Integer.parseInt(st.nextToken());
        boolean flag = true;
        for(int i=1; i<M; i++){
            if(find(target) != find(Integer.parseInt(st.nextToken())))
                flag = false;
        }

        System.out.println(flag ? "YES" : "NO");

모두 같은 집합에 있어야만 모든 도시들을 여행할 수 있기 때문에 flag를 두어서 같은 집합에 속하지 않는 도시가 있다면 NO를 출력한다.

📌 전체 코드

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

class Main {
    static int[] A;
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

        int N = Integer.parseInt(br.readLine())+1;
        int M = Integer.parseInt(br.readLine());
        A = new int[N];

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

        StringTokenizer st;
        for(int i=1; i<N; i++){
            st = new StringTokenizer(br.readLine(), " ");
            for(int j=1; j<N; j++){
                boolean c = st.nextToken().equals("1") ? true : false;
                if(c) union(i, j);
            }
        }

        st = new StringTokenizer(br.readLine());
        int target = Integer.parseInt(st.nextToken());
        boolean flag = true;
        for(int i=1; i<M; i++){
            if(find(target) != find(Integer.parseInt(st.nextToken())))
                flag = false;
        }

        System.out.println(flag ? "YES" : "NO");

    }

    static void union(int a, int b){
        int x = find(a);
        int y = find(b);

        A[y] = x;
    }
    static int find(int a){
        if(A[a] == a){
            return a;
        }else{
            A[a] = find(A[a]);
            return A[a];
        }
    }

}

결과