백준 28251 - 나도리합

문제 링크: https://acmicpc.net/problem/28251 

28251번: 나도리합

---

문제를 요약해봅시다.

• 각 나도리는 각각 어떤 그룹 $G_x$에 속해 있으며, 그룹 $G_x$의 전투력은 $G_x$에 속한 서로 다른 두 나도리의 크기의 곱의 총합, 즉 $\sum_{i, j \in G_x, i < j}s_i s_j$입니다.
• 쿼리 $a\\;b$가 들어오면 나도리 $a$와 $b$가 속한 그룹을 합쳐 하나의 그룹으로 만들고, 합쳐진 그룹의 전투력을 출력해야 합니다.
• 문제에 정확히 나와 있지는 않는 것 같지만, 모든 나도리는 처음에 자기 자신만이 속한 크기 1인 그룹에 속해 있습니다.

---

각 나도리는 반드시 하나의 그룹에만 속해야 하고, 두 그룹을 합칠 수 있다고 하니 분리 집합으로 풀어봅시다.

최초에 모든 나도리는 자기 자신만이 속한 그룹에 있으므로 모든 그룹의 전투력은 $0$입니다.

이후 쿼리 $a\;b$가 들어오면, $a$가 속한 그룹과 $b$가 속한 그룹을 합친 후 전투력을 계산하여 출력한 후 저장하면 됩니다.

 

이제 전투력을 계산해봅시다. 그룹 \(X\)와 그룹 \(Y\)를 합치는 상황을 생각해 봅시다.

\(X\)의 크기를 \(p\)라 하고, 여기에 속한 나도리를 각각 \(x_1, x_2, x_3, ... , x_p\)라 합시다. 비슷하게 \(Y\)의 크기를 \(q\)라 하고, 여기에 속한 나도리를 각각 \(y_1, y_2, y_3, ... , y_q\)라 합시다.

전투력을 구하려면 각각의 \(x_i\)와 \(y_j\)를 모두 곱해서 더해보면 되겠지만.. 나도리의 수가 \(N\leq 200,000\)으로 상당히 크므로 그룹을 합칠 때마다 전부 계산하는 것은 시간초과가 날 것입니다. 계산을 줄여봅시다.

 

그룹의 전투력은 그룹 내 서로 다른 두 나도리의 크기의 곱을 모두 합한 것입니다. \(X\)와 \(Y\)를 합친 새로운 그룹의 전투력을 구하는데는 다음 세 가지 종류의 나도리 쌍을 생각해야 합니다.

• 두 나도리 모두 원래 \(X\)에 속해 있던 나도리인 경우.
• 두 나도리 모두 원래 \(Y\)에 속해 있던 나도리인 경우.
• 두 나도리 중 하나는 \(X\)에, 하나는 \(Y\)에 속해 있던 나도리인 경우.

 

두 나도리가 모두 \(X\)에 속해 있던 경우, 이 나도리 쌍들의 크기의 곱을 모두 더한 것은 \(X\)의 전투력과 같다는 것을 알 수 있습니다. 비슷하게, 둘 다 \(Y\)에 속해 있던 경우는 \(Y\)의 전투력과 같다는 것을 알 수 있습니다.

따라서 세 번째 경우인 하나는 \(X\)에 속해 있었고, 나머지 하나는 \(Y\)에 속해 있었던 경우를 생각해봅시다.

 

\(X\)에 속해 있던 나도리 중 \(Y\)에도 속해 있던 나도리는 없습니다. 왜냐하면 각 나도리는 하나의 그룹에만 들어 있기 때문입니다.

그렇다면, \(X\)의 나도리의 크기 \(s_{x_i}\)에 대해 \(y_1, y_2, ... , y_q\)의 크기를 각각 곱하고, 모두 더해주면 됩니다.

식으로 표현하면 다음과 같습니다.

\[\sum_{i=1}^{p}s_{x_i}\times\sum_{j=1}^{q}s_{y_j}\]

즉, 원래 \(X\)에 속해 있던 나도리의 크기의 합과 \(Y\)에 속해 있던 나도리의 크기의 합을 곱해주면 되는겁니다!

 

따라서 새로운 그룹의 전투력은 다음 식으로 표현할 수 있게 됩니다. 그룹 \(X, Y\)의 전투력을 각각 \(P_X, P_Y\)라 하면,

\[P_{X \cup Y} = P_X + P_Y + \sum_{i=1}^{p}s_{x_i}\times\sum_{j=1}^{q}s_{y_j}\]

입니다. 이를 쉽게 구현하려면 각 그룹의 전투력 뿐 아니라, 그룹에 속한 나도리들의 크기 합도 저장해두면 됩니다.

 

시간 복잡도는 다음과 같습니다.

• 크기 \(O(N)\)인 배열을 분리 집합 연산을 위해 초기화하므로 \(O(N)\)
• 분리 집합 연산은 아커먼 역함수에 비례하는 시간 복잡도를 갖고, 이를 \(Q\)번 반복하므로 \(O(Q\alpha(N))\)

 

아커먼 역함수가 굉장히 작아 사실 \(O(1)\)로 생각한다면, 총합 \(O(N + Q)\)의 시간 복잡도를 갖게 됩니다.

---

소스코드 (C++)

#include <bits/stdc++.h>

#define fastio ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);

using namespace std;
typedef long long ll;

// 그룹의 크기 합을 저장
vector<ll> set_sum;
// 그룹의 전투력을 저장
vector<ll> set_power;
vector<ll> parent;

// 그룹 검색 연산
int find_parent(int v) {
    if (v == parent[v]) return v;
    else return parent[v] = find_parent(parent[v]);
}

// 그룹 병합 연산
void union_parent(int v1, int v2) {
    int p1 = find_parent(v1);
    int p2 = find_parent(v2);
    
    if (p1 > p2) swap(p1, p2);
    
    if (p1 != p2) {
        parent[p1] = p2;
        set_power[p2] = set_power[p1] + set_power[p2] + set_sum[p1] * set_sum[p2];
        set_sum[p2] = set_sum[p1] + set_sum[p2];
    }
}

int main() {
    fastio;
    
    int n, q;
    cin >> n >> q;
    
    set_sum = vector<ll>(n + 1);
    set_power = vector<ll>(n + 1);
    parent = vector<ll>(n + 1);
    
    for (int i = 1; i <= n; i++) {
        cin >> set_sum[i];
        parent[i] = i;
    }
    
    while (q--) {
        int a, b;
        cin >> a >> b;

        union_parent(a, b);
        cout << set_power[find_parent(a)] << "\n";
    }
}