けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AOJ 3180 GCDMST (HUPC 2020 day3-I)

中盤枠という感じで作られた

問題へのリンク

問題概要

 1, \dots, N の番号を振られた  N 個の頂点があります。 最初、これらを繋ぐ辺はありません。 あなたはいくつかの辺を追加してこのグラフを連結にしたいと思いました。

頂点  i j を繋ぐ辺を追加するには  A_{{\rm gcd}(i, j)} のコストがかかります。 このグラフを連結にするように辺を追加するとき、かかるコストの和の最小値を求めてください。

制約

  •  2 \le N \le 10^{5}

考えたこと

普通に kruskal 法でできそう!

  • A が小さい順に処理していく
  • w = A[v] として、辺 (v, 2v), (v, 3v), ... それぞれについて連結でないならば重さ w で繋ぐ

という風にすれば OK。このとき、考える辺の本数は「調和級数の和」の形になるので、 O(N \log N) となる。よって全体の計算量は  O(N \alpha(N) \log N)

コード

#include <bits/stdc++.h>
using namespace std;
template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; }
template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return 1; } return 0; }

struct UnionFind {
    vector<int> par;
    
    UnionFind(int n) : par(n, -1) { }
    void init(int n) { par.assign(n, -1); }
    
    int root(int x) {
        if (par[x] < 0) return x;
        else return par[x] = root(par[x]);
    }
    
    bool issame(int x, int y) {
        return root(x) == root(y);
    }
    
    bool merge(int x, int y) {
        x = root(x); y = root(y);
        if (x == y) return false;
        if (par[x] > par[y]) swap(x, y); // merge technique
        par[x] += par[y];
        par[y] = x;
        return true;
    }
    
    int size(int x) {
        return -par[root(x)];
    }
};

using pll = pair<long long, int>;
int main() {
    int N;
    cin >> N;
    vector<pll> A(N);
    for (int i = 0; i < N; ++i) cin >> A[i].first, A[i].second = i + 1;
    sort(A.begin(), A.end());
    UnionFind uf(N);
    long long res = 0;
    for (int i = 0; i < N; ++i) {
        long long w = A[i].first;
        int u = A[i].second;
        for (int v = u * 2; v <= N; v += u) {
            if (uf.issame(u-1, v-1)) continue;
            res += w;
            uf.merge(u-1, v-1);
        }
    }
    cout << res << endl;
}