[백준/Gold 4] 여행가자

문제 정보

여행가자 - 1976

시간 제한: 2초

메모리 제한: 128MB

문제 설명

동혁이는 친구들과 함께 여행을 가려고 한다. 한국에는 도시가 N개 있고 임의의 두 도시 사이에 길이 있을 수도, 없을 수도 있다. 동혁이의 여행 일정이 주어졌을 때, 이 여행 경로가 가능한 것인지 알아보자. 물론 중간에 다른 도시를 경유해서 여행을 할 수도 있다.

예를 들어 도시가 5개 있고, A-B, B-C, A-D, B-D, E-A의 길이 있고, 동혁이의 여행 계획이 E C B C D 라면 E-A-B-C-B-C-B-D라는 여행경로를 통해 목적을 달성할 수 있다.

도시들의 개수와 도시들 간의 연결 여부가 주어져 있고, 동혁이의 여행 계획에 속한 도시들이 순서대로 주어졌을 때 가능한지 여부를 판별하는 프로그램을 작성하시오. 같은 도시를 여러 번 방문하는 것도 가능하다.

입력

• 첫 줄에 도시의 수 N이 주어진다. (N ≤ 200)
• 둘째 줄에 여행 계획에 속한 도시들의 수 M이 주어진다. (M ≤ 1000)
• 다음 N개의 줄에는 N개의 정수가 주어진다. i번째 줄의 j번째 수는 i번 도시와 j번 도시의 연결 정보를 의미한다.
  • 1이면 연결된 것
  • 0이면 연결되지 않은 것
• A와 B가 연결되었으면 B와 A도 연결되어 있다.
• 마지막 줄에는 여행 계획이 주어진다.
• 도시의 번호는 1부터 N까지 차례대로 매겨져 있다.

출력

첫 줄에 가능하면 YES, 불가능하면 NO를 출력한다.

풀이

위 문제는 Union-Find를 문제로 보인다.

각 도시를 (0부터 시작한다고 가정) parents[] 라는 배열을 생성해 root 를 저장하는 방식으로 도시의 연결을 체크할려고 한다.

초기 상태에 모든 도시의 root는 자기 자신을 가진다.

1. 초기 root 설정하기

parents = IntStream.rangeClosed(0, N - 1).toArray(); // 도시들의 부모 기록

• IntStream.rangeClosed(0, N - 1); : 0부터 N-1까지의 정수 Stream을 생성
• .toArray() Stream을 배열로 변환

2. 도시의 연결을 설정하기

int root = parent(city); // 도시의 root를 기록
for (int dest = 0; dest < N; dest++) { // 해당 도시의 목적지의 root를 찾아서 갱신  
    if (Integer.parseInt(st.nextToken()) == 1 && root != parent(dest)) {  
        parents[parent(dest)] = root;  
    }  
}

N번째 줄마다 입력을 받고 해당 입력이 1일 경우 양 도시의 root를 비교한다. 만약 root가 다르다면 목적지인 parents[parent(dest)]에 N 번째 도시를 root로 설정한다.

3. 방문가능한지 확인

boolean ans = Arrays.stream(visited).allMatch(i -> parent(i) == parent(visited[0]));
System.out.println(ans ? "YES" : "NO");

방문이 가능한지 확인할 도시 배열 visited[]를 선언 후 .allMatch를 통해 i 번째 index의 도시의 부모와 가장 최초의 visited[0]의 부모가 모두 같은지 비교한다.

전체 코드

import java.io.BufferedReader;  
import java.io.InputStreamReader;  
import java.util.*;  
import java.util.stream.IntStream;  
  
public class Main {  
    static int N, M;  
    static int[] parents;  
    public static void main(String[] args) throws Exception {  
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));  
        StringTokenizer st;  
  
        N = Integer.parseInt(br.readLine());  
        M = Integer.parseInt(br.readLine());  
  
        parents = IntStream.rangeClosed(0, N - 1).toArray(); // 도시들의 부모 기록  
  
        for (int city = 0; city < N; city++) {  
            st = new  StringTokenizer(br.readLine());  
  
            int root = parent(city); // 도시의 root를 기록  
            for (int dest = 0; dest < N; dest++) { // 해당 도시의 목적지의 root를 찾아서 갱신  
                if (Integer.parseInt(st.nextToken()) == 1 && city < dest && root != parent(dest)) {  
                    parents[parent(dest)] = root;  
                }  
            }  
        }  
  
        st = new StringTokenizer(br.readLine());  
        int[] visited = new int[M];  
        for (int i = 0; i < M; i++) {  
            visited[i] = Integer.parseInt(st.nextToken()) - 1;  
        }  
  
        boolean ans = Arrays.stream(visited).allMatch(i -> parent(i) == parent(visited[0]));  
        System.out.println(ans ? "YES" : "NO");  
    }  
  
    public static int parent(int city) {  
        if (parents[city] == city) {  
            return city;  
        }  
  
        return parents[city] = parent(parents[city]); // 경로 압축  
    }  
}