DMOPC '15 Contest 6 P4 - Line Graph

    1#include <bits/stdc++.h>
    2using namespace std;
    3typedef long long ll;
    4const int INF = 0x3f3f3f3f; //❄️
    5struct DSU {
    6    int n;
    7    vector<int> parent, rank;
    8    DSU (int _n) : n(_n) , parent(n + 1), rank(n + 1, 0) {
    9        for (int i = 0; i < parent.size(); i++) parent[i] = i;
    10    }
    11    void unite(int x, int y) {
    12        x = find(x); y = find(y);
    13        if (rank[x] > rank[y]) {
    14            parent[y] = x;
    15        } else if (rank[y] > rank[x]) {
    16            parent[x] = y;
    17        } else {
    18            parent[x] = y;
    19            rank[y]++;
    20        }
    21    }
    22    int find(int x) {
    23        if(parent[x] == x) return x;
    24        return parent[x] = find(parent[x]);
    25    }
    26};
    27struct edge {
    28    int u, v, w;
    29    bool operator<(const edge& other) const{
    30        return w < other.w;
    31    }
    32};
    33int main() {
    34    ios::sync_with_stdio(0); cin.tie(0); 
    35    int n, k; cin >> n >> k;
    36    vector<edge> v;
    37    for (int i = 1; i <= n; i++) {
    38        int w; cin >> w;
    39        if(i == n) break;
    40        v.push_back({i, i + 1, w});
    41    }
    42    for(int i = 1; i <= n; i++){
    43        if (i + k > n) {
    44            break;
    45        }
    46        v.push_back({i, i + k, 0});
    47    }
    48    sort(v.begin(), v.end());
    49    DSU dsu(n);
    50    int ans = 0, cnt = 0;
    51    for (int i = 0; i < v.size(); i++) {
    52        auto[x, y, w] = v[i];
    53        if (dsu.find(x) != dsu.find(y)) {
    54            dsu.unite(x, y);
    55            ans += w; cnt++;
    56        }
    57        if(cnt == n - 1) break;
    58    }
    59    cout << ans << '\n';
    60    return 0;
    61}