섹션 9 Greedy - 6. 친구인가?(Disjoint-Set : Union&Find)

문제 설명

현수는 반 친구들의 친구관계를 알고 싶어 한다. 각 학생은 1~N번까지 번호가 붙어 있으며 친구 쌍이 주어질 때, 두 학생이 직접 또는 간접적으로 친구인지를 판별해야 한다.

• 예: (1,2), (2,3), (3,4) → 1과 4는 친구

---

입력 & 출력

입력 예시

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

 

출력 예시

NO

---

해결 전략: Union-Find

• 친구 관계는 서로소 집합 (Disjoint Set) 구조로 표현할 수 있다.
• 같은 집합에 속하면 친구, 아니라면 친구 아님.
• 각 학생은 부모 노드를 통해 루트(대표 친구)를 찾고,
  서로 다른 루트를 가진 두 학생을 친구로 연결(합침)한다.

---

전체 코드

import java.util.*;

public class Problem6 {
	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 sc = new Scanner(System.in);
		int n = sc.nextInt();
		int m = sc.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 = sc.nextInt();
			int b = sc.nextInt();
			Union(a, b);
		}

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

		sc.close();
	}
}

---

코드 설명

• static int[] unf: Union-Find를 위한 부모 배열 (unf[i]는 i의 대표 번호)
• Find(int v)
  • 노드 v의 대표 노드(루트)를 찾음
  • 경로 압축 기법 적용 → unf[v] = Find(unf[v])
  • 시간 복잡도 거의 O(1)
• Union(int a, int b)
  • 두 원소 a, b의 루트를 찾아서 하나로 합침
  • 이미 같은 집합이면 아무 일도 하지 않음
  • unf[fa] = fb를 통해 병합
• main()
  • 입력값: 학생 수 n, 친구 관계 개수 m
  • 초기화: unf[i] = i로 자기 자신이 루트가 되도록 설정
  • m개의 친구 관계를 입력받아 Union(a, b)로 병합
  • 마지막 줄의 두 학생 a, b에 대해 Find(a) == Find(b)인지 판단
    → 같으면 "YES", 아니면 "NO" 출력

---

시간 복잡도

• Find, Union 모두 거의 O(1) (경로 압축 포함 시)
• 입력 크기 N ≤ 1,000, M ≤ 3,000 → 매우 빠르게 처리 가능

---

출처: 인프런 자바(Java) 알고리즘 문제풀이 입문: 코딩테스트 대비

https://www.inflearn.com/course/자바-알고리즘-문제풀이-코테대비/dashboard

자바(Java) 알고리즘 문제풀이 입문: 코딩테스트 대비 강의 | 김태원 - 인프런