CF #597 (Div. 2) D. Shichikuji and Power Grid

Problem

Problem - D - Codeforces

Solution

・Excluding the building power plant, it is clear that getting MST is answer.
・finding that which power plant is efficient is difficult. (e.g. greedy algo is not correct.like this) also, if we build power plants on and , these are not needed to connect each other .
To solve this problem, Let's consider new vertex as building power plant.

・ is new vertex, connecting to means build power plants at with cost .
・Note that we have to build power plants at least one. therefor, getting MST which include new vertex is answer.
・Time complexity is .

#include<bits/stdc++.h>
using namespace std;
using ll = long long
#define int ll
class unionfind {
    vector<int> par, rank, size_;
public:
    unionfind(int n) :par(n), rank(n), size_(n, 1) {
        iota(par.begin(),par.end(), 0);
    }
    int find(int x) {
        if (par[x] == x)return x;
        return par[x] = find(par[x]);
    }
    void unite(int x, int y) {
        x = find(x), y = find(y);
        if (x == y)return;
        if (rank[x] < rank[y])swap(x, y);
        par[y] = x;
        size_[x] += size_[y];
        if (rank[x] == rank[y])rank[x]++;
    }
    bool same(int x, int y) {
        return find(x) == find(y);
    }
    int size(int x) {
        return size_[find(x)];
    }
};
auto dis(pair<int,int> t, pair<int,int> u){
    return abs(t.first - u.first) + abs(t.second - u.second);
}
signed main(){
    int N; cin >> N;
    vector<pair<int, int>>pos(N);
    for(auto&  p : pos)
        cin >> p.first  >> p.second;
    vector<int>c(N), k(N);
    for(auto& e : c)  cin >> e;
    for(auto& e : k)  cin >> e;
    vector<vector<pair<int,int>>>edge(N);
    using tp = tuple<int,int,int>;
    vector<tp>g;
    for(int i = 0; i < edge.size(); ++i){
        auto t = pos[i];
        for(int j = 0; j < N; ++j){
            if(i == j)continue;
            auto u = pos[j];
            int cost = dis(t,u)*(k[i] + k[j]);
            edge[i].emplace_back(j,cost);
            g.emplace_back(cost,i,j);
        }
    }
    int over_v = N + 1;
    for(int i = 0; i < N; ++i)
        g.emplace_back(c[i], over_v, i);
    
    sort(g.begin(),g.end());
    auto uf = unionfind((int)g.size() + 10);
    int ans = 0;
    vector<bool>build(N,false);
    vector<pair<int,int>>used_edge;
    for(auto p : g){
        int cost, t,u;
        tie(cost,t,u) = p;
        if(uf.same(t,u)) continue;
        ans += cost;
        uf.unite(t,u);
        if(t == over_v or u == over_v){
            build[(t == over_v ? u : t)] = true;
        } else{
            used_edge.emplace_back(t,u);
        }
    }
    cout << ans << endl;
    cout << count_if(build.begin(),build.end(),[](bool e){ return e;}) << endl;;
    int cb = 0;
    for(int i = 0; i < build.size(); ++i){
        if(build[i]){
            if(cb++) cout << " ";
            cout << i + 1;
        }
    }
    cout << endl;
    cout << used_edge.size() << endl;
    for(auto p : used_edge)
        cout << p.first + 1 << " " << p.second + 1 << endl;