Solution for "Intel Code Challenge Elimination Round (Div. 1 + Div. 2, combined)" Problem C, Destruction of Array

cf_intel_elem.cpp

/*
Solution for "Intel Code Challenge Elimination Round (Div. 1 + Div. 2, combined)"
Problem C, Destruction of Array
http://codeforces.com/contest/722/problem/C

I used modified DisjointSet (additionally stores sum of elements in connected component)
Bitmask to check neighbors
*/
#include <bits/stdc++.h>
using namespace std;

class disjoint_sets
{
    vector<int> parent;
    vector<int> rank;
    vector<long long> sum;
public:
    disjoint_sets(size_t sz) : parent(sz, -1), rank(sz, -1), sum(sz, 0) {}

    void make_set(int v, int weight)
    {
        parent[v] = v;
        rank[v] = 0;
        sum[v] = weight;
    }

    int find_set(int x)
    {
        if (parent[x] != x)
            return parent[x] = find_set(parent[x]);
        return x;
    }

    long long sum_set(int x) {
        return sum[find_set(x)];
    }

    void union_sets(int l, int r)
    {
        int pl = find_set(l);
        int pr = find_set(r);
        if (pl != pr) {
            if (rank[pl] < rank[pr])
                swap(pl, pr);
            parent[pr] = pl;
            if (rank[pr] == rank[pl])
                rank[pl] += 1;
            sum[pl] += sum[pr];
        }
    }
};

vector<int> a, d;
void solve() {
    disjoint_sets ds(a.size());
    vector<bool> neigh(a.size(), false);
    stack<long long> res;
    res.push(0);
    long long maxRootSum = 0;
    for (int i = d.size() - 1; i > 0; --i) {
        int pos = d[i] - 1;
        ds.make_set(pos, a[pos]);

        neigh[pos] = true;
        if (pos > 0 && neigh[pos-1]) {
            ds.union_sets(pos, pos-1);
        }
        if (pos < a.size() - 1 && neigh[pos+1]) {
            ds.union_sets(pos, pos+1);
        }
        maxRootSum = max(maxRootSum, ds.sum_set(pos));
        
        res.push( maxRootSum );
    }

    while (!res.empty()) {
        cout << res.top() << endl;
        res.pop();
    }
}

int main(int, char**) {
    std::ios::sync_with_stdio(false);
    int n;
    cin >> n;
    a.resize(n);
    d.resize(n);

    for (int i = 0; i < n; ++i)
        cin >> a[i];

    for (int i = 0; i < n; ++i)
        cin >> d[i];
    
    solve();

    return 0;
}