백준 1976번(여행가자)[JAVA]

풀이

유니온 파인드 문제이다. 유니온 파인드(Union-find)란, 여러 노드가 있을 때 특정 2개의 노드를 연결하여 하나의 집합을 묶는 union연산, 특정 노드가 속한 집합의 대표 노드를 찾는 find연산으로 이루어져 있다.
union연산은 인덱스와 값이 같게 초기화된 배열에서 union된 두 노드 중 하나 인덱스의 값을 다른 노드 인덱스의 값(일명 대표 노드)으로 바꿔주면 끝이다. find는 특정 노드의 대표 노드를 찾으면 바로 그 값을 리턴하고, 아니라면 대표 노드를 재귀 호출로 계속 찾아낸다.
이 문제는 경로가 행렬로 주어졌으므로, 행렬을 이중 for문으로 탐색하며 union-find연산을 해주면 된다.

1. 행렬 값이 1인 곳의 인덱스(i, j)를 union연산을 해준다.
2. 여행 경로의 첫번째를 집합의 대표 노드로 한다.
3. 여행 경로의 두번째부터 find연산을 하여 집합의 대표 노드와 같은지 확인하고, 이 결과에 따라 출력을 달리한다.

소스 코드

import java.util.*;

public class 여행가자 {
	static int[] arr;
	public static void main(String[] args) {
		Scanner sc = new Scanner(System.in);
		int n = sc.nextInt();
		int m = sc.nextInt();
		arr = new int[n+1];
		for (int i = 1; i < n+1; i++) {
			arr[i] = i;
		}
		
		int[][] city = new int[n+1][n+1];
		for (int i = 1; i < n+1; i++) {
			for (int j = 1; j < n+1; j++) {
				city[i][j] = sc.nextInt();
			}
		}
		
		int[] route = new int[m+1];
		for (int i = 1; i < m+1; i++) {
			route[i] = sc.nextInt();
		}
		
		for (int i = 1; i < n+1; i++) {
			for (int j = 1; j < n+1; j++) {
				if (city[i][j] == 1) 
					union(i,j);
			}
		}
		
		int idx = find(route[1]);
		for (int i = 2; i < route.length; i++) {
			if (idx != find(route[i])) {
				System.out.println("NO");
				return;
			}
		}
		System.out.println("YES");
	}
	
	static void union(int a, int b) {
		a = find(a);
		b = find(b);
		if (a != b)
			arr[b] = a;
	}
	
	static int find(int a) {
		if (a == arr[a])
			return a;
		else
			return arr[a] = find(arr[a]);
	}
}

태클 감사합니다.
조언 환영입니다.