백준 3830번 C++ 풀이

문제

https://www.acmicpc.net/problem/3830

시간 제한	메모리 제한	solved.ac 티어
2초	256MB	플래티넘 3

사진을 클릭하면 문제로 이동합니다.

 

---

풀이

 이 문제는 풀고 기분이 아주 좋았다. 2일 반동안 풀이에 대해 고민했고, 구현을 구체화하며, 그 과정에서 알고리즘에 대한 이해가 조금 오른, 정말 성장했다는 기분을 느끼게 해주는 AC였다. 오랜만에 느끼는 희열이다.

 

 원래 골드 이하 문제를 많이 풀고 있던 나였기에, 이 문제를 보고 처음에는 그래프 최단거리 문제로 접근하고자 했다. 그리고 나름의 최적화를 위해 각 샘플들을 유니온 파인드에 저장하여, 같은 집합이 아니라면 최단거리 탐색을 하지 않도록, 개선하고자 했다. 그래프 최단거리 알고리즘 세가지를 모두 따져보자.

• 플로이드-워셜 : 시간복잡도 O(N^3)에 공간복잡도 O(N^2)로 TLE MLE 둘다 발생한다.
• 다익스트라 : 시간복잡도의 측면에서 다익스트라 1회에는 O(MlogN)이고, 이를 M번의 쿼리에 대해 적용해야한다. 애매하게 될듯 말듯하는데, 다익스트라를 못쓰는 이유가 있다. 해당 그래프에는 음수간선이 등장할 수 있다.
• 벨만포드 : a->b가 w라면, b->a는 -w이기 때문에 음수간선이 등장할 수 있다. 해당 그래프는 방향그래프이지만, 무방향그래프처럼 두 노드를 자유롭게 왕복할 수 있어야한다. 그렇기에 다익스트라가 불가능하며, 벨만포드 알고리즘은 1회에 O(NM)의 시간복잡도를 가지고, 당연히 TLE가 발생할 것이다.

이런 이유로 그래프 최단거리 알고리즘을 사용하는 문제가 아님을 알 수 있다. 그러다 문득, 유니온 파인드에 무게차이를 저장해보면 어떨까 하는 아이디어가 떠올랐다. 그리고 각 무게차이를 모두 집합의 루트와의 무게차이로 변경하면 두번의 find연산으로 처리가 가능할 것이라는 생각이 들었다. a와 b의 무게차이 b-a는 a와 root와의 무게차이 A = root - a, b와 root와의 무게차이 B = root - b를 가지고 B-A를 통해 구할수 있기 때문이다. A와 B는 유니온 파인드를 업데이트하면서 만들어줘야 한다.

 

 유니온 파인드를 이용해서 어떻게 푸는지 구체적으로 알아보도록 하자. 어떤 두 샘플의 무게차이를 구할 수 있다면, 두 샘플은 같은 집합에 속해있다고 볼 수 있다. 예를 들어서 1, 2와 2와 3의 무게차이 정보가 있고, 5, 6끼리의 무게차이 정보가 있다면, 1과 6의 무게차이를 알 수 있는가? 불가능하다. 다른 집합이기 때문이다. 하지만 1과 3의 무게차이는 알 수 있을까? 가능하다. 1, 2, 3은 2라는 샘플로 인해 1과 3의 무게차이를 이끌어낼 수 있다.

 

 그렇기에 ? 쿼리가 들어온다면, 우선 같은 집합인지부터 알아낸다. 같은 집합이 아니라면 무게차이를 구할 수 없기 때문에 UNKNOWN을 출력, 같은 집합이라면 각 샘플과 집합의 루트와의 무게차이를 구한 후, 빼준다. 이때 우리가 구해야하는건 b가 a보다 얼마나 무거운지, 즉 b-a가 얼마인지를 구해야한다. 하지만 루트와의 무게는 각각 r-a, r-b이기 때문에 (루트와 a의 무게차이) - (루트와 b의 무게차이)를 구하면 된다. 해당 무게차이는 위에서 말했듯 유니온 파인드의 내부 배열에 데이터가 저장되고 갱신된다.

 

 ! 쿼리가 들어와 무게차이 정보가 들어온다면, 두 샘플을 merge하고, 무게차이가 w가 되도록 업데이트 한다. 이때, 두 샘플은 각기 다른 집합에 속해있을수도 있고, 둘다 집합의 크기가 1인 (다른 샘플과 무게차이 데이터가 없는) 집합일수도 있다. 후자의 경우라면 간선의 데이터를 w로 저장하면 되지만, 후자의 경우는 유니온 파인드의 특성상 루트와 루트끼리 연결되기 때문에 루트와 루트가 연결되었을 때, 무게의 차이가 w가 되도록 두 간선의 값을 결정해야 한다. ! a b w라고 했을 때, a와 b의 루트를 r_a, r_b라고 하자. cost[a], cost[b]는 각각 r_a - a, r_b - b의 값을 의미할 것이다. 이때 r_a와 r_b가 merge되기 때문에, 우리는 r_b - r_a를 b - a가 w가 되도록 설정해야 한다(merge(a, b)를 할 때, b 기준으로 합쳐지도록 merge한다고 하자.) 최종적으로 w라는 값이 되도록 하려면

 

(기존 b의 cost) - (기존 a의 코스트) + w = (r_b - b) - (r_a - a) + (b - a) = r_b - r_a

 

라는 식을 쓰면 된다. 만약 둘다 루트가 본인인(다른 무게정보가 없으면) a와 b의 코스트가 0이었을 테니 w가 저장된다.

 

코드

#include <bits/stdc++.h>
#define ALL(X) X.begin(), X.end()
#define endl '\n'
using namespace std;
using ll = long long;

typedef struct Disjoint_set{
    int* tree;
    ll* cost;

    Disjoint_set(int n) {
        tree = new int[n + 1], cost = new ll[n + 1];
        memset(tree, 0xff, sizeof(int) * (n + 1));
        memset(cost, 0, sizeof(ll) * (n + 1));
    }

    void merge(int a, int b, int w) {
        int ra = find(a), rb = find(b); // find를 먼저 수행하여, 간선의 정보를 모두 업데이트해준다.
        ll ac = cost[a], bc = cost[b];
        if (ra == rb) return;

        tree[rb] += tree[ra];
        tree[ra] = rb;

        // (r_b - b) - (r_a - a) + b - a = r_b - r_a 를 구할 수 있다. 만약 두 노드 모두 연결되어있지 않던 노드라면 bc, ac 모두 0이기 때문에 w가 저장된다.
        cost[ra] = bc - ac + w; 
    }

    int find(int a, int prev = -1) {
        if (tree[a] < 0) {
            return a;
        }
        tree[a] = find(tree[a], a);

        if (prev > 0) { // 이전 노드에 간선 비용 누적합
            cost[prev] += cost[a];
        }

        return tree[a];
    }

    ~Disjoint_set() { delete[] tree; }
} DS;

int main() {
    cin.tie(0)->sync_with_stdio(false);
    int n, m, a, b, w;
    char q;

    while(1) {
        cin >> n >> m;
        if (!n && !m) break;
        DS ds(n + 1);

        
        while(m--) {
            cin >> q >> a >> b;
            if (q == '!') {
                cin >> w;
                ds.merge(a, b, w);
            }
            else {
                if (ds.find(a) != ds.find(b)) {
                    cout << "UNKNOWN" << endl;
                }
                else {
                    cout << ds.cost[a] - ds.cost[b] << endl; // (r - a) - (r - b) = b - a
                }
            }
        }
    }
}