백준 1976 여행 가자 (C++)

1976번: 여행 가자

오랜만에 분리 집합 문제이다. 먼저 root를 초기화해주고 반복문을 돌며 도시가 연결되있을 경우 root를 union 해준다. 그 후 여행 경로에 있는 도시들의 root를 비교해 결과를 출력해주었다. 오랜만에 분리 집합 문제라 시간이 오래 걸렸다. 그리고 문제를 제대로 읽지 않아 배열 범위를 잘못 적어 오류가 나기도 했다. 문제를 잘 읽어보자.

#include <iostream>

using namespace std;

int N, M;
int A[200][200];
int B[1000];
int root[20];
string result = "YES";

int findRoot(int a) {
    if (root[a] == a) return a;
    else return root[a] = findRoot(root[a]);
}

void unionRoot(int a, int b) {
    a = findRoot(a);
    b = findRoot(b);

    if (a != b) root[a] = b;
}

void solution() {
    for (int i = 0; i < N; i++) {
        root[i] = i;
    }

    for (int i = 0; i < N; i++) {
        for (int j = 0; j < N; j++) {
            if (A[i][j]) {
                unionRoot(i, j);
            }
        }
    }

    for (int i = 0; i < M - 1; i++) {
        if (findRoot(B[i] - 1) != findRoot(B[i + 1] - 1)) {
            result = "NO";
            break;
        }
    }

    cout << result;
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL); cout.tie(NULL);

    cin >> N;
    cin >> M;

    for (int i = 0; i < N; i++) {
        for (int j = 0; j < N; j++) {
            cin >> A[i][j];
        }
    }

    for (int i = 0; i < M; i++) {
        cin >> B[i];
    }

    solution();

    return 0;
}