Java : Disjoint-Set (Union-Find)

문제

풀이

• 유니온 파인드 문제
• union(합집합 연산) : 두 원소 a,b가 있으면 각 원소가 속한 집합을 하나로 합친다.
• find(찾기 연산) : 어떠한 원소 a가 주어지면 해당 원소가 속한 집합(최상위 노드)을 찾아서 반환
• 즉 학생 관계를 a,b라고 했을 때 Union으로 학생들이 속한 집합(=친구관계)을 구하고 출력 결과인 친구가 맞는지(find(a)==find(b))를 판단하면 된다.

잡담

[1,2,3,4,5] 의 학생이 있다고 치자
여기서 1,2를 하나의 집합으로(=친구로!) 한다고 하면
[1,1,3,4,5] 가 된다. 2가 1에 속한 것이다.

만약 (1,2), (1,4), (2,3), (4,5) 이렇게 친구 관계라면

누구든지 find(a)를 호출하면 부모 노드를 찾아가기 때문에 1을 반환하게 된다.(전부 친구)

숫자는 일반적으로 작은 수 밑으로 들어간다
유니온 파인드 참고 사이트

전체 코드

package inflearn;

import java.util.Scanner;

public class I0906 {
  static int[] parent;

  public static void main(String[] args) {

    Scanner sc = new Scanner(System.in);
    int n = sc.nextInt();
    int m = sc.nextInt();
    parent = new int[n + 1];

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

    for (int i = 0; i < m; i++) {
      int a = sc.nextInt();
      int b = sc.nextInt();
      union(a, b);
    }

    int s1 = sc.nextInt();
    int s2 = sc.nextInt();

    System.out.println(find(s1) == find(s2) ? "YES" : "NO");
  }

  static int find(int s1) {
    if (parent[s1] == s1) return s1;
    else return parent[s1] = find(parent[s1]);
  }

  static void union(int s1, int s2) {
    s1 = find(s1);
    s2 = find(s2);

    if (s1 == s2) return;

    if (s1 < s2) parent[s2] = s1;
    else parent[s1] = s2;
  }

}