ARC029 C問題 高橋君と国家

グラフの問題。最小全域木かなぁと思ったけど、都市に対してもコストがあったので、無理だなぁと重い最初は愚直に書くもTLE。
ある適当な仮想国家を一つ持ってきて、そこから都市への道をその都市のコストにしてあげると綺麗になって、最終的にプリム法で一発。
この考え方は結構大事だなぁ。

struct edge{
    int from, to, cost;
    edge(){}
    edge(int f, int t, int c):from(f), to(t), cost(c){}
};

const int MAX = 100001;
int N, M;
int C[MAX];
edge G[2*MAX];
P CG[MAX*3];

struct UnionFind
{
    std::vector<int> data;
    UnionFind(int size) : data(size, -1){}
    void initialize(void){
        for(int i=0; i<(int)data.size(); i++) data[i] = i;
    }
    bool merge(int x, int y){
        x = find(x); y = find(y);
        if(x == y) return false;
        else{ data[x] = y; return true; }
    }
    int find(int x){ //根っこを見つける関数
        if(data[x] == x) return x;
        else return data[x] = find(data[x]); //経路圧縮
    }
    bool isSame(int x, int y){
        return find(x) == find(y);
    }
};

int main()
{
    cin.tie(0);
    ios::sync_with_stdio(false);

    cin >> N >> M;
    for(int i=0; i<N; i++){
        cin >> C[i];
        CG[i] = P(C[i], i);
    }
    for(int i=0; i<M; i++){
        int a, b, r;
        cin >> a >> b >> r; a--; b--;
        G[i] = edge(a, b, r);
        CG[i+N] = P(r, N+i);
    }
    sort(CG, CG+N+M);

    UnionFind uf(N+M+1);
    uf.initialize();
    ll res = 0;
    for(int i=0; i<N+M; i++){
        int index = CG[i].second;
        if(index < N){
            if(!uf.isSame(index, N+M)){
                uf.merge(index, N+M);
                res += CG[i].first;
            }
        }else{
            if(!uf.isSame(G[index-N].to, G[index-N].from)){
                uf.merge(G[index-N].to, G[index-N].from);
                res += CG[i].first;
            }
        }
    }
    cout << res << endl;
}