Problem p03143 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <iostream>
#include <queue>
#include <vector>

using namespace std;

struct GroupState {
  int maxei;
  long long vwsum;
};

struct UnionFind {
  vector<int> par;
  vector<long long> sum;
  vector<int> last_e;

  UnionFind(int n, vector<int> &vw) {
    par.resize(n);
    sum.resize(n);
    last_e.resize(n, -1);
    for (int i = 0; i < n; ++i) {
      par[i] = i;
      sum[i] = vw[i];
    }
  }

  int find(int x) {
    if (x == par[x])
      return x;
    return par[x] = find(par[x]);
  }

  pair<GroupState, GroupState> unite(int x, int y, int ei) {
    int xr = find(x), yr = find(y);
    auto ret = make_pair(GroupState{last_e[xr], sum[xr]},
                         GroupState{last_e[yr], sum[yr]});

    long long xsum = sum[xr];
    par[find(x)] = yr;
    sum[yr] += xsum;
    last_e[yr] = ei;

    return ret;
  }

  bool same(int x, int y) { return find(x) == find(y); }
};

struct Edge {
  int src, dst, cost, index;
};
bool operator<(const Edge &e, const Edge &f) { return e.cost < f.cost; }

int main() {
  int n, m;
  cin >> n >> m;
  vector<int> vw(n);
  vector<Edge> edge(m);

  vector<pair<GroupState, GroupState>> to_check(m);

  for (int i = 0; i < n; ++i)
    cin >> vw[i];
  for (int i = 0; i < m; ++i) {
    int a, b, c;
    cin >> a >> b >> c;
    edge[i] = {a - 1, b - 1, c, i};
  }

  auto sorted_edge = edge;
  sort(sorted_edge.begin(), sorted_edge.end());
  vector<long long> prev_vw(m), prev_e(m);

  UnionFind uf(n, vw);
  vector<int> uf_edge(m);
  for (auto &e : sorted_edge) {
    int x = e.src, y = e.dst;
    if (uf.same(x, y))
      continue;
    auto res = uf.unite(x, y, e.index);
    to_check[e.index] = res;
    uf_edge[e.index] = 1;
    //
    //        cerr << "new e" << ' ' << e.index << endl;
    //        cerr << "new e cost" << ' ' << e.cost << endl;
    //        cerr << "new sum" << ' ' << uf.sum[uf.find(x)] << endl;
    //        cerr << "prev e1: " << to_check[e.index].first.maxei << endl;
    //        cerr << "prev sum1: " << to_check[e.index].first.vwsum << endl;
    //        cerr << "prev e2: " << to_check[e.index].second.maxei << endl;
    //        cerr << "prev sum2: " << to_check[e.index].second.vwsum << endl;
    //        cerr << "---" << endl;
  }

  int root = uf.find(0);
  cerr << root << endl;
  int ei = uf.last_e[root];
  long long sum = uf.sum[root];

  vector<int> can_use(m, 1);

  queue<GroupState> q;
  q.push(GroupState{ei, sum});
  while (!q.empty()) {
    auto gs = q.front();
    q.pop();
    int idx = gs.maxei;

    //        cerr << "Edge index: " << idx << endl;
    //        cerr << "Edge cost: " << edge[idx].cost<< endl;
    //        cerr << "Group weight: " << gs.vwsum << endl;

    if (edge[idx].cost <= gs.vwsum) {
      continue;
    }
    auto pgs = to_check[idx];
    if (pgs.first.maxei >= 0) {
      q.push(pgs.first);
      //            cerr << "push " << pgs.first.maxei << endl;
    }
    if (pgs.second.maxei >= 0) {
      //            cerr << "push " << pgs.second.maxei << endl;
      q.push(pgs.second);
    }
    can_use[idx] = 0;
  }

  int edgecnt = 0;
  UnionFind uf2(n, vw);
  for (auto &e : sorted_edge) {
    if (!uf_edge[e.index] || !can_use[e.index])
      continue;
    int x = e.src, y = e.dst;
    if (uf2.same(x, y))
      continue;
    auto res = uf2.unite(x, y, e.index);
    to_check[e.index] = res;
    uf_edge[e.index] = 1;
    ++edgecnt;
  }

  cerr << "rest edge" << endl;

  for (int i = 0; i < m; ++i) {
    if (uf_edge[i])
      continue;
    if (!uf2.same(edge[i].src, edge[i].dst))
      continue;
    int root = uf2.find(edge[i].src);
    if (uf2.sum[root] >= edge[i].cost) {
      cerr << i << endl;
      ++edgecnt;
    }
  }

  cout << m - edgecnt << endl;
}

Fixed submission:
#include <algorithm>
#include <iostream>
#include <queue>
#include <vector>

using namespace std;

struct GroupState {
  int maxei;
  long long vwsum;
};

struct UnionFind {
  vector<int> par;
  vector<long long> sum;
  vector<int> last_e;

  UnionFind(int n, vector<int> &vw) {
    par.resize(n);
    sum.resize(n);
    last_e.resize(n, -1);
    for (int i = 0; i < n; ++i) {
      par[i] = i;
      sum[i] = vw[i];
    }
  }

  int find(int x) {
    if (x == par[x])
      return x;
    return par[x] = find(par[x]);
  }

  pair<GroupState, GroupState> unite(int x, int y, int ei) {
    int xr = find(x), yr = find(y);
    auto ret = make_pair(GroupState{last_e[xr], sum[xr]},
                         GroupState{last_e[yr], sum[yr]});

    long long xsum = sum[xr];
    par[find(x)] = yr;
    sum[yr] += xsum;
    last_e[yr] = ei;

    return ret;
  }

  bool same(int x, int y) { return find(x) == find(y); }
};

struct Edge {
  int src, dst, cost, index;
};
bool operator<(const Edge &e, const Edge &f) { return e.cost < f.cost; }

int main() {
  int n, m;
  cin >> n >> m;
  vector<int> vw(n);
  vector<Edge> edge(m);

  vector<pair<GroupState, GroupState>> to_check(m);

  for (int i = 0; i < n; ++i)
    cin >> vw[i];
  for (int i = 0; i < m; ++i) {
    int a, b, c;
    cin >> a >> b >> c;
    edge[i] = {a - 1, b - 1, c, i};
  }

  auto sorted_edge = edge;
  sort(sorted_edge.begin(), sorted_edge.end());
  vector<long long> prev_vw(m), prev_e(m);

  UnionFind uf(n, vw);
  vector<int> uf_edge(m);
  for (auto &e : sorted_edge) {
    int x = e.src, y = e.dst;
    if (uf.same(x, y))
      continue;
    auto res = uf.unite(x, y, e.index);
    to_check[e.index] = res;
    uf_edge[e.index] = 1;
    //
    //        cerr << "new e" << ' ' << e.index << endl;
    //        cerr << "new e cost" << ' ' << e.cost << endl;
    //        cerr << "new sum" << ' ' << uf.sum[uf.find(x)] << endl;
    //        cerr << "prev e1: " << to_check[e.index].first.maxei << endl;
    //        cerr << "prev sum1: " << to_check[e.index].first.vwsum << endl;
    //        cerr << "prev e2: " << to_check[e.index].second.maxei << endl;
    //        cerr << "prev sum2: " << to_check[e.index].second.vwsum << endl;
    //        cerr << "---" << endl;
  }

  int root = uf.find(0);
  cerr << root << endl;
  int ei = uf.last_e[root];
  long long sum = uf.sum[root];

  vector<int> can_use(m, 1);

  queue<GroupState> q;
  if (ei >= 0)
    q.push(GroupState{ei, sum});
  while (!q.empty()) {
    auto gs = q.front();
    q.pop();
    int idx = gs.maxei;

    //        cerr << "Edge index: " << idx << endl;
    //        cerr << "Edge cost: " << edge[idx].cost<< endl;
    //        cerr << "Group weight: " << gs.vwsum << endl;

    if (edge[idx].cost <= gs.vwsum) {
      continue;
    }
    auto pgs = to_check[idx];
    if (pgs.first.maxei >= 0) {
      q.push(pgs.first);
      //            cerr << "push " << pgs.first.maxei << endl;
    }
    if (pgs.second.maxei >= 0) {
      //            cerr << "push " << pgs.second.maxei << endl;
      q.push(pgs.second);
    }
    can_use[idx] = 0;
  }

  int edgecnt = 0;
  UnionFind uf2(n, vw);
  for (auto &e : sorted_edge) {
    if (!uf_edge[e.index] || !can_use[e.index])
      continue;
    int x = e.src, y = e.dst;
    if (uf2.same(x, y))
      continue;
    auto res = uf2.unite(x, y, e.index);
    to_check[e.index] = res;
    uf_edge[e.index] = 1;
    ++edgecnt;
  }

  cerr << "rest edge" << endl;

  for (int i = 0; i < m; ++i) {
    if (uf_edge[i])
      continue;
    if (!uf2.same(edge[i].src, edge[i].dst))
      continue;
    int root = uf2.find(edge[i].src);
    if (uf2.sum[root] >= edge[i].cost) {
      cerr << i << endl;
      ++edgecnt;
    }
  }

  cout << m - edgecnt << endl;
}