Problem p03143 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <bitset>
#include <cmath>
#include <cstring>
#include <ext/algorithm>
#include <ext/numeric>
#include <iostream>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <vector>
#define ffor(_a, _f, _t) for (int _a = (_f), __t = (_t); _a < __t; _a++)
#define all(_v) (_v).begin(), (_v).end()
#define sz size()
#define pb push_back
#define SET(__set, val) memset(__set, val, sizeof(__set))
#define FOR(__i, __n) ffor(__i, 0, __n)
typedef long long LL;
using namespace std;

const int MAXN = 400000;

const int UNION_FIND_N = MAXN;
int parent[UNION_FIND_N];
LL sum[UNION_FIND_N];
int lastEdge[UNION_FIND_N];
int vWeight[MAXN];

vector<pair<int, pair<int, int>>> edges;

bool goodEdge[MAXN];
LL sumAtGood[MAXN];
vector<int> G[MAXN];

// assumes that if vertex v does not have parent,
// i.e. if v is root, then parent[v] = v
int root(int v) {
  if (parent[v] == v)
    return v;
  return parent[v] = root(parent[v]);
}

void merge(int u, int v, int edge) {
  int rv = root(v);
  int ru = root(u);
  if (rv == ru)
    return;

  parent[rv] = ru;
  sum[ru] += sum[rv];
  G[edge].pb(lastEdge[rv]);
  G[edge].pb(lastEdge[ru]);
  lastEdge[ru] = edge;
  if (edges[edge].first <= sum[ru]) {
    goodEdge[edge] = true;
    sumAtGood[edge] = sum[ru];
  }
}

void init_union_find(int n) {
  for (int i = 0; i < n; i++) {
    parent[i] = i;
    sum[i] = vWeight[i];
    lastEdge[i] = -i - 1;
  }
}

int n, m;

int label[MAXN];

void dfs(int edge, bool good, int l) {
  bool nextGood = good;
  int nextLabel = l;
  if (!good && goodEdge[edge]) {
    nextGood = true;
    nextLabel = edge;
  }

  int e0 = G[edge][0], e1 = G[edge][1];

  if (e0 < 0) {
    e0 = -e0 - 1;
    label[e0] = nextLabel;
  } else {
    dfs(e0, nextGood, nextLabel);
  }

  if (e1 < 0) {
    e1 = -e1 - 1;
    label[e1] = nextLabel;
  } else {
    dfs(e1, nextGood, nextLabel);
  }
}

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

  FOR(i, MAXN)
  G[i].clear();
  edges.clear();

  cin >> n >> m;
  FOR(i, n)
  cin >> vWeight[i];

  int x, y, c;
  FOR(i, m) {
    cin >> x >> y >> c;
    x--;
    y--;
    edges.pb(make_pair(c, make_pair(x, y)));
  }
  sort(all(edges));

  init_union_find(n);

  SET(goodEdge, 0);

  int lastAdded = -1;

  FOR(i, edges.sz) {
    int a = edges[i].second.first;
    int b = edges[i].second.second;
    if (root(a) == root(b))
      continue;

    merge(a, b, i);
    lastAdded = i;
  }

  dfs(lastAdded, false, -1);

  int ret = 0;
  FOR(i, edges.sz) {
    int a = edges[i].second.first;
    int b = edges[i].second.second;
    if (label[a] < 0 || label[b] < 0) {
      ret++;
      continue;
    }
    if (label[a] != label[b]) {
      ret++;
      continue;
    }
    if (sumAtGood[label[a]] < edges[i].first)
      ret++;
  }

  cout << ret << endl;
  return 0;
}

Fixed submission:
#include <algorithm>
#include <bitset>
#include <cmath>
#include <cstring>
#include <ext/algorithm>
#include <ext/numeric>
#include <iostream>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <vector>
#define ffor(_a, _f, _t) for (int _a = (_f), __t = (_t); _a < __t; _a++)
#define all(_v) (_v).begin(), (_v).end()
#define sz size()
#define pb push_back
#define SET(__set, val) memset(__set, val, sizeof(__set))
#define FOR(__i, __n) ffor(__i, 0, __n)
typedef long long LL;
using namespace std;

const int MAXN = 400000;

const int UNION_FIND_N = MAXN;
int parent[UNION_FIND_N];
LL sum[UNION_FIND_N];
int lastEdge[UNION_FIND_N];
int vWeight[MAXN];

vector<pair<int, pair<int, int>>> edges;

bool goodEdge[MAXN];
LL sumAtGood[MAXN];
vector<int> G[MAXN];

// assumes that if vertex v does not have parent,
// i.e. if v is root, then parent[v] = v
int root(int v) {
  if (parent[v] == v)
    return v;
  return parent[v] = root(parent[v]);
}

void merge(int u, int v, int edge) {
  int rv = root(v);
  int ru = root(u);
  if (rv == ru)
    return;

  parent[rv] = ru;
  sum[ru] += sum[rv];
  G[edge].pb(lastEdge[rv]);
  G[edge].pb(lastEdge[ru]);
  lastEdge[ru] = edge;
  if (edges[edge].first <= sum[ru]) {
    goodEdge[edge] = true;
    sumAtGood[edge] = sum[ru];
  }
}

void init_union_find(int n) {
  for (int i = 0; i < n; i++) {
    parent[i] = i;
    sum[i] = vWeight[i];
    lastEdge[i] = -i - 1;
  }
}

int n, m;

int label[MAXN];

void dfs(int edge, bool good, int l) {
  bool nextGood = good;
  int nextLabel = l;
  if (!good && goodEdge[edge]) {
    nextGood = true;
    nextLabel = edge;
  }

  int e0 = G[edge][0], e1 = G[edge][1];

  if (e0 < 0) {
    e0 = -e0 - 1;
    label[e0] = nextLabel;
  } else {
    dfs(e0, nextGood, nextLabel);
  }

  if (e1 < 0) {
    e1 = -e1 - 1;
    label[e1] = nextLabel;
  } else {
    dfs(e1, nextGood, nextLabel);
  }
}

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

  FOR(i, MAXN)
  G[i].clear();
  edges.clear();

  cin >> n >> m;
  FOR(i, n)
  cin >> vWeight[i];

  if (!m) {
    cout << 0 << endl;
    return 0;
  }

  int x, y, c;
  FOR(i, m) {
    cin >> x >> y >> c;
    x--;
    y--;
    edges.pb(make_pair(c, make_pair(x, y)));
  }
  sort(all(edges));

  init_union_find(n);

  SET(goodEdge, 0);

  int lastAdded = -1;

  FOR(i, edges.sz) {
    int a = edges[i].second.first;
    int b = edges[i].second.second;
    if (root(a) == root(b))
      continue;

    merge(a, b, i);
    lastAdded = i;
  }

  dfs(lastAdded, false, -1);

  int ret = 0;
  FOR(i, edges.sz) {
    int a = edges[i].second.first;
    int b = edges[i].second.second;
    if (label[a] < 0 || label[b] < 0) {
      ret++;
      continue;
    }
    if (label[a] != label[b]) {
      ret++;
      continue;
    }
    if (sumAtGood[label[a]] < edges[i].first)
      ret++;
  }

  cout << ret << endl;
  return 0;
}