친구인가

설명

오늘은 새 학기 새로운 반에서 처음 시작하는 날이다. 현수네 반 학생은 N명이다. 현수는 각 학생들의 친구관계를 알고 싶다.

모든 학생은 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개의 줄에 걸쳐 숫자쌍이 주어진다.

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

 

출력

첫 번째 줄에 “YES"또는 "NO"를 출력한다.

 

예시 입력 1 

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

 

예시 출력 1

NO

 

문제 해결

 

이 문제는 서로소 집합(Union-Find) 알고리즘을 사용해 친구 관계를 확인하는 문제이다.

 

 

서로소 집합(Union-Find)

 

서로소 집합 알고리즘은 여러 개의 원소가 있을 때,

서로소 부분 집합들로 나누고, 두 원소가 같은 집합에 속하는지를 판단하는 자료구조

 

Find: 특정 원소가 어떤 집합에 속하는지 찾는 연산

Union: 두 집합을 합치는 연산

 

 

 

즉 Union & Find 연산의 흐름을 이번 문제에 적용시켜보면

• 학생 수 N과 친구 관계의 수 M을 입력받고, 각 친구 관계를 입력받아 Union 연산을 수행한 뒤
• 주어진 두 학생이 같은 집합에 속하는지 Find 연산으로 확인하는 과정으로 이루어진다 !!

 

<Find 연산>

public static int Find(int v) {
    if (v == unf[v]) return v;
    else return unf[v] = Find(unf[v]);
}

• v가 자기 자신을 가리키고 있다면, v는 자기 자신의 루트
• 그렇지 않은 경우, v의 부모(unf[v])에 대해 Find 연산을 재귀적으로 호출. 이때 반환된 값을 unf[v]에 저장해 v를 루트에 직접 연결 (경로 압축).

 

 

<Union 연산>

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

• 두 집합을 하나의 집합으로 합치는 연산!
• a의 루트와 b의 루트를 찾아 두 원소의 루트가 다르다면, 한 집합의 루트를 다른 집합의 루트로 설정해 두 집합을 합치는 작업을 해준다. 

 

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;

• 학생 수 n과 친구 관계 수 m을 입력받고 인덱스 번호와 일치시기게 하기 위해 unf 배열 크기를 n+1로 설정한다.
• 그리고 모든 학생은 자기 자신을 집합의 대표로 설정하도록 초기화시킨다. (처음에는 모든 학생이 독립적인 집합을 가지고 있음)

 

for (int i = 1; i <= m; i++) {
    int a = in.nextInt();
    int b = in.nextInt();
    Union(a, b);
}

• 주어진 M개의 친구 관계에 대해 Union 연산을 수행해 두 학생이 같은 집합에 속하도록 한다.

int a = in.nextInt();
int b = in.nextInt();
int fa = Find(a);
int fb = Find(b);
if (fa == fb) System.out.println("YES");
else System.out.println("NO");

• 마지막으로 입력된 두 학생의 대표를 Find 연산으로 찾아, 두 학생이 같은 집합에 속하는지 확인한다
  • 만약 두 학생의 대표가 같다면 "YES",
  • 그렇지 않다면 "NO"를 출력한다

 

최종 코드

 

import java.util.Scanner;

public class P06_친구인가_Uion_Find알고리즘 {
    static int[] unf;

    public static int Find(int v) {
        if (v == unf[v]) return 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 static void main(String[] args) {
        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 = 1; i <= m; i++) {
            int a = in.nextInt();
            int b = in.nextInt();
            Union(a, b);
        }

        int a = in.nextInt();
        int b = in.nextInt();
        int fa = Find(a);
        int fb = Find(b);
        if (fa == fb) System.out.println("YES");
        else System.out.println("NO");
    }
}