[JAVA] Union-Find

Union-Find 란?

• Union : 특정 2개의 노드를 연결해 1개의 집합으로 묶는 작업
• Find : 특정 2개의 노드가 같은 집합에 속해 있는지 확인하는 작업

• 다음과 같이 6개의 노드가 있다고 가정
• (1, 2), (2, 3), (4, 6), (3, 6) 을 연결했을 때, 생기는 그룹은 몇개일까?
  1. group[1~6] = [1, 2, 3, 4, 5, 6]으로 초기화
     • 이는 각 index가 속한 그룹의 대표값이라 생각하면 됨
  2. 1, 2 연결
  3. 노드 1과 노드 2의 대표값이 다름 => group[2]에 1을 넣어줌 => group = [1, 1, 3, 4, 5, 6]
  4. 2, 3 연결
  5. 노드2와 노드 3의 대표값이 다름 => group[3]에 노드2의 대표값 1을 넣어줌 => group = [1, 1, 1, 4, 5, 6]
  6. 4, 6 연결
  7. 노드 4와 노드 6의 대표값이 다름 => group[4]에 6을 넣어줌 => group = [1, 1, 1, 4, 5, 4]
  8. 3, 6 연결
  9. 이 둘을 연결하면 1, 2, 3, 4, 6이 모두 같은 그룹에 들어가게 됨
  10. 노드 3의 대표값을 찾음 = 1
  11. 노드 6의 대표값을 찾음 = 4
  12. group[4]에 1을 넣어줌 => group[1, 1, 1, 1, 5, 4]
  13. 노드 1, 2, 3, 4, 6을 따라가보면 대표값은 모두 1로 같다는 것을 알 수 있음
      • 모두 같은 그룹에 속함
• 시간 단축을 하려면 group[6]의 값을 1로 바꿔주는 작업이 필요
• 나중에 6의 대표값을 찾게 되는 타이밍에 재귀함수를 통해 1을 찾고 돌아오면서 갱신

구현

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

1717번: 집합의 표현

 

• 이 문제를 Union-Find를 통해 해결하면서 구현해봄

코드

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

public class p1717 {
    static int[] rep;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st = new StringTokenizer(br.readLine());
        int n = Integer.parseInt(st.nextToken());
        int m = Integer.parseInt(st.nextToken());

        rep = new int[n + 2];
        for(int i = 0 ; i <= n ; i ++) {
            rep[i] = i;
        }

        StringBuilder sb = new StringBuilder();
        for(int i = 0 ; i < m ; i ++) {
            st = new StringTokenizer(br.readLine());
            int op = Integer.parseInt(st.nextToken());
            int a = Integer.parseInt(st.nextToken());
            int b = Integer.parseInt(st.nextToken());

            int aRep = getRep(a);
            int bRep = getRep(b);

            // 합집합 연산
            if(op == 0) {
                rep[aRep] = bRep;
            }
            // 두 원소가 같은 집합인지 판단
            else {
                if(aRep == bRep) {
                    sb.append("YES").append("\n");
                } else {
                    sb.append("NO").append("\n");
                }
            }
        }

        System.out.print(sb);
    }

    // 대표값 구하기
    static int getRep(int index) {
        if(rep[index] == index) {
            return index;
        } else {
            rep[index] = getRep(rep[index]);
            return rep[index];
        }
    }
}