Computer >> 컴퓨터 >  >> 프로그래밍 >> C++

C++로 가중치 그래프에서 참이 되는 쿼리 개수 계산하기 (유니온-파인드)

문제 개요

무방향 그래프의 간선 목록이 주어져 있다고 가정해 보겠습니다. 각 간선은 [u, v, w] 형태로 표현되며, 여기서 u는 출발 정점, v는 도착 정점, w는 해당 간선의 가중치를 의미합니다. 이와 함께 동일한 형식 [u, v, w]를 가지는 쿼리 목록도 주어집니다. 각 쿼리는 "정점 u에서 정점 v로 가는 경로가 존재하며, 경로에 포함된 모든 간선의 가중치가 w 이하인 경우가 있는가?"라는 질문을 나타냅니다. 우리의 목표는 이러한 쿼리 중 참(true)인 것의 개수를 구하는 것입니다.

예를 들어 입력이 다음과 같다고 합시다.

edges = [[0, 1, 6], [1, 2, 7], [2, 3, 8], [0, 3, 5]], queries = [[0, 2, 14], [1, 0, 3]]

C++로 가중치 그래프에서 참이 되는 쿼리 개수 계산하기 (유니온-파인드)

그렇다면 출력은 1이 됩니다. 첫 번째 쿼리의 경우, 노드 0에서 노드 2까지 경로 [0, 1, 2]를 따라 이동할 수 있으며, 이 경로의 간선 가중치는 각각 6과 7로 모두 14 이하입니다(총 가중치는 13). 반면 두 번째 쿼리는 노드 1에서 노드 0으로 이동하는 질문인데, 가중치 3 이하의 간선만으로는 어떤 경로로도 도달할 수 없으므로 거짓입니다.

접근 방법: 오프라인 처리와 유니온-파인드

이 문제는 간선과 쿼리를 하나로 묶어 가중치 기준으로 정렬한 뒤 순서대로 처리하는 '오프라인(offline)' 기법으로 효율적으로 해결할 수 있습니다. 핵심 아이디어는 다음과 같습니다.

  • 간선과 쿼리를 한 배열에 넣고 가중치 오름차순으로 정렬합니다. 이때 간선과 쿼리를 구분하기 위한 타입 플래그(간선은 0, 쿼리는 1)를 함께 저장합니다.
  • 배열을 앞에서부터 순회하면서 간선을 만나면 유니온(Union) 연산으로 두 정점을 같은 집합으로 합칩니다.
  • 쿼리를 만나는 시점에는 이미 가중치가 해당 쿼리의 임계값 w 이하인 간선들이 모두 병합된 상태입니다. 따라서 두 정점이 같은 집합에 속해 있는지, 즉 루트 부모가 같은지만 확인하면 됩니다. 같다면 그 쿼리는 참입니다.

전체 시간 복잡도는 정렬 단계가 지배하므로 O((E + Q) log(E + Q))이며, 경로 압축(path compression)을 적용한 유니온-파인드 연산은 사실상 상수 시간에 처리됩니다.

알고리즘 단계

이 문제를 해결하기 위해 다음 단계를 따릅니다.

  1. get_parent(x, par) 함수를 정의합니다. par[x]가 자기 자신이 아니면 재귀적으로 루트를 찾아 par[x]에 다시 저장합니다(경로 압축). 그리고 par[x]를 반환합니다.
  2. 메인 solve() 메서드에서 2차원 배열 gr을 선언하고, n := 0으로 초기화합니다.
  3. edges의 각 간선 t에 대해: n을 max(n, t[0], t[1])로 갱신하고, gr에 행 [t[2], 0, t[0], t[1]](가중치, 타입 0=간선, 정점 a, 정점 b)을 삽입합니다.
  4. queries의 각 쿼리 t에 대해: gr에 행 [t[2], 1, t[0], t[1]](타입 1=쿼리)을 삽입합니다.
  5. gr을 오름차순으로 정렬합니다.
  6. 크기가 n + 1인 배열 par를 만들고, i := 0부터 n까지 par[i] := i로 초기화합니다.
  7. sz := queries의 크기, ans := 0으로 설정합니다.
  8. gr의 각 행 t에 대해:
    • a := t[2], b := t[3], tp := t[1], d := t[0]
    • px := get_parent(a, par), py := get_parent(b, par)
    • tp가 0(간선)이면: px ≠ py일 때 par[py] := px로 두 집합을 병합합니다.
    • tp가 1(쿼리)이면: px == py이면 ans를 1 증가시킵니다. 이후 sz를 1 감소시키고, sz가 0이 되면 모든 쿼리를 처리했으므로 반복을 종료합니다.
  9. ans를 반환합니다.

C++ 구현 예제

더 나은 이해를 돕기 위해 다음 구현을 살펴보겠습니다.

#include <bits/stdc++.h>
using namespace std;
int get_parent(int x, vector<int>& par) {
    if (par[x] != x) {
        par[x] = get_parent(par[x], par);
    }
    return par[x];
}
int solve(vector<vector<int>>& edges, vector<vector<int>>& queries) {
    vector<vector<int>> gr;
    int n = 0;
    for (auto t : edges) {
        n = max(n, max(t[0], t[1]));
        gr.push_back({t[2], 0, t[0], t[1]});
    }
    for (auto t : queries) {
        gr.push_back({t[2], 1, t[0], t[1]});
    }
    sort(gr.begin(), gr.end());
    vector<int> par(n + 1, -1);
    for (int i = 0; i <= n; i++) {
        par[i] = i;
    }
    int sz = queries.size();
    int ans = 0;
    for (auto t : gr) {
        int a = t[2];
        int b = t[3];
        int tp = t[1];
        int d = t[0];
        int px = get_parent(a, par);
        int py = get_parent(b, par);
        if (tp == 0) {
            if (px != py) {
                par[py] = px;
            }
        }else {
            if (px == py) {
                ans++;
            }
            sz--;
            if(sz == 0) {
                break;
            }
        }
    }
    return ans;
}
int main(){
    vector<vector<int>> edges = {{0, 1, 6},{1, 2, 7},{2, 3, 8},{0, 3, 5}};
    vector<vector<int>> queries = {{0, 2, 14},{1, 0, 3}};
    cout << solve(edges, queries);
}

입력

{{0, 1, 6},{1, 2, 7},{2, 3, 8},{0, 3, 5}}, {{0, 2, 14},{1, 0, 3}}

출력

1

마무리

이처럼 쿼리를 미리 정렬해두고 간선을 가중치 순서대로 병합하면서 답을 기록하는 오프라인 기법은, 매 쿼리마다 그래프를 새로 탐색하는 비효율적인 방법(O(Q × (E + V)))보다 훨씬 빠릅니다. 유니온-파인드 자료구조와 결합하면 대규모 그래프에서도 안정적인 성능을 보장할 수 있으므로, 코딩 테스트나 경쟁 프로그래밍에서 자주 활용되는 패턴이니 꼭 익혀두시길 바랍니다.