본문 바로가기
알고리즘 이전/그리디

친구인가(Disjoint-set: 유니온-파인드)

by hoshi03 2023. 9. 6.

유니온 파인드 만드는건 외워두자

 

모든 학생은 1부터 N까지 번호가 부여되어 있고, 현수에게는 각각 두 명의 학생은 친구 관계가 번호로 표현된 숫자쌍이 주어진다.

만약 (1, 2), (2, 3), (3, 4)의 숫자쌍이 주어지면 1번 학생과 2번 학생이 친구이고, 2번 학생과 3번 학생이 친구, 3번 학생과 4번 학생이 친구이다.

그리고 1번 학생과 4번 학생은 2번과 3번을 통해서 친구관계가 된다.

학생의 친구관계를 나타내는 숫자쌍이 주어지면 특정 두 명이 친구인지를 판별하는 프로그램을 작성하세요.

두 학생이 친구이면 “YES"이고, 아니면 ”NO"를 출력한다.

첫 번째 줄에 반 학생수인 자연수 N(1<=N<=1,000)과 숫자쌍의 개수인 M(1<=M<=3,000)이 주어지고,

다음 M개의 줄에 걸쳐 숫자쌍이 주어진다.

마지막 줄에는 두 학생이 친구인지 확인하는 숫자쌍이 주어진다.

 

입력

9 7 
1 2
2 3
3 4
1 5
6 7
7 8
8 9
3 8

출력

NO

 

 

유니온파인드로 서로소 인지 판단하는 문제, 유니온파인드는 각 인덱스를 값으로 가지는 unf배열, union, find 함수가 필요하다

find함수는 v의 값이 unf[v]와 다르면 재귀적으로 find를 호출해서 부모 노드의 값을 찾는다

find 함수에서 unf[v] = find(unf[v]);를 해주면서 그래프를 압축시키면 시간 복잡도가 좋아지니 return unf[v] = find(unf[v]);를 하는걸 기억하자

union 함수는 인자로 들어온 두 자연수 a,b로 find를 호출하고 값이 다르면  unf[fa]에 fb값을 덮어씌워서 같은 집합이 되게 만든다

 static int[] unf;

    public static int find(int v){
        if (v == unf[v]) return v;
        //unf[v]의 값을 저장해두면서 그래프를 압축
        else return unf[v] = find(unf[v]);
    }


    public static void union(int a, int b){
        int fa = find(a);
        int fb = find(b);
        if (fa != fb) unf[fa] = fb;
    }

전체 코드

 

import java.util.*;

class Main{

    static int[] unf;

    public static int find(int v){
        if (v == unf[v]) return v;
        //unf[v]의 값을 저장해두면서 그래프를 압축
        else return unf[v] = find(unf[v]);
    }


    public static void union(int a, int b){
        int fa = find(a);
        int fb = find(b);
        if (fa != fb) unf[fa] = fb;
    }

    public void Solution(){


    }

    public static void main(String[] args) {
        Main T = new Main();
        Scanner in = new Scanner(System.in);
        int n = in.nextInt();
        int m = in.nextInt();
        unf = new int[n+1];
        for (int i = 1; i <= n; i++) unf[i] = i;
        for (int i =0; i < m; i++){
            int a = in.nextInt();
            int b = in.nextInt();
            union(a, b);
        }
        int a = in.nextInt();
        int b = in.nextInt();
        if (find(a) == find(b)) System.out.println("YES");
        else System.out.println("NO");

    }
}