Problem p03143 (C++)

Buggy submission (status: Runtime Error):
#include <algorithm>
#include <cassert>
#include <cctype>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <tuple>
#include <typeinfo>
#include <unordered_map>
#include <unordered_set>
#include <vector>
using namespace std;

#define REP2(i, m, n) for (int i = (int)(m); i < (int)(n); i++)
#define REPD(i, n) for (int i = (int)(n)-1; i >= 0; i--)
#define REP(i, n) REP2(i, 0, n)
#define ALL(c) (c).begin(), (c).end()
#define PB(e) push_back(e)
#define FOREACH(i, c) for (auto i = (c).begin(); i != (c).end(); ++i)
#define MP(a, b) make_pair(a, b)
#define BIT(n, m) (((n) >> (m)) & 1)

typedef long long ll;

template <typename S, typename T>
ostream &operator<<(ostream &out, const pair<S, T> &p) {
  out << "(" << p.first << ", " << p.second << ")";
  return out;
}

template <typename T> ostream &operator<<(ostream &out, const vector<T> &v) {
  out << "[";
  REP(i, v.size()) {
    if (i > 0)
      out << ", ";
    out << v[i];
  }
  out << "]";
  return out;
}

class UnionFind {
public:
  int num_components;
  std::vector<int> parent;
  std::vector<ll> weight;
  std::vector<int> rank;

public:
  UnionFind(int N, const vector<ll> weight)
      : num_components(N), parent(std::vector<int>(N)), weight(weight),
        rank(std::vector<int>(N, 0)) {
    for (int i = 0; i < N; i++)
      parent[i] = i;
  }

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

  int size(int x) { return weight[find(x)]; }

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

  void unite(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return;

    num_components--;
    if (rank[x] < rank[y]) {
      weight[y] += weight[x];
      parent[x] = y;
    } else {
      weight[x] += weight[y];
      parent[y] = x;
      if (rank[x] == rank[y])
        rank[y]++;
    }
  }

  int count() { return num_components; }
};

int main(int argc, char *argv[]) {
  int N, M;
  cin >> N >> M;
  vector<ll> WN(N);
  REP(i, N) cin >> WN[i];
  vector<pair<ll, pair<int, int>>> E(M);
  REP(i, M) {
    int a, b, y;
    cin >> a >> b >> y;
    a--;
    b--;
    E[i] = make_pair(y, make_pair(a, b));
  }
  sort(ALL(E));
  UnionFind uf(N, WN);

  vector<tuple<ll, ll, ll>> op;
  vector<int> valid(N - 1, true);
  vector<int> last_op(N, -1);
  vector<vector<int>> ch_ops(N - 1);

  REP(i, M) {
    int w = E[i].first;
    // cout << w << endl;
    int a = E[i].second.first;
    int b = E[i].second.second;
    if (uf.same(a, b)) {
      continue;
    }
    // cout << "merge: " << a << " " << b << endl;
    a = uf.find(a);
    b = uf.find(b);
    int c = op.size();
    if (last_op[a] != -1) {
      ch_ops[c].push_back(last_op[a]);
    }
    if (last_op[b] != -1) {
      ch_ops[c].push_back(last_op[b]);
    }
    op.push_back(make_tuple(uf.weight[a], uf.weight[b], i));
    uf.unite(a, b);
    last_op[uf.find(a)] = c;
  }
  assert(int(op.size()) == N - 1);
  queue<int> que;
  que.push(N - 2);
  while (!que.empty()) {
    int i = que.front();
    que.pop();
    ll s0 = get<0>(op[i]);
    ll s1 = get<1>(op[i]);
    ll w = E[get<2>(op[i])].first;
    // cout << i << " " << w << " " << s0 << " "<< s1 << endl;
    if (w > s0 + s1) {
      valid[i] = false;
      for (int j : ch_ops[i]) {
        que.push(j);
      }
    }
  }
  UnionFind new_uf = UnionFind(N, WN);
  REP(i, N - 1) {
    if (valid[i]) {
      int eid = get<2>(op[i]);
      int a = E[eid].second.first;
      int b = E[eid].second.second;
      new_uf.unite(a, b);
    }
  }
  ll res = 0;
  REP(i, M) {
    int w = E[i].first;
    int a = new_uf.find(E[i].second.first);
    int b = new_uf.find(E[i].second.second);
    if (!new_uf.same(a, b) || w > new_uf.weight[a]) {
      res++;
    }
  }
  cout << res << endl;
  return 0;
}


Fixed submission:
#include <algorithm>
#include <cassert>
#include <cctype>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <tuple>
#include <typeinfo>
#include <unordered_map>
#include <unordered_set>
#include <vector>
using namespace std;

#define REP2(i, m, n) for (int i = (int)(m); i < (int)(n); i++)
#define REPD(i, n) for (int i = (int)(n)-1; i >= 0; i--)
#define REP(i, n) REP2(i, 0, n)
#define ALL(c) (c).begin(), (c).end()
#define PB(e) push_back(e)
#define FOREACH(i, c) for (auto i = (c).begin(); i != (c).end(); ++i)
#define MP(a, b) make_pair(a, b)
#define BIT(n, m) (((n) >> (m)) & 1)

typedef long long ll;

template <typename S, typename T>
ostream &operator<<(ostream &out, const pair<S, T> &p) {
  out << "(" << p.first << ", " << p.second << ")";
  return out;
}

template <typename T> ostream &operator<<(ostream &out, const vector<T> &v) {
  out << "[";
  REP(i, v.size()) {
    if (i > 0)
      out << ", ";
    out << v[i];
  }
  out << "]";
  return out;
}

class UnionFind {
public:
  int num_components;
  std::vector<int> parent;
  std::vector<ll> weight;
  std::vector<int> rank;

public:
  UnionFind(int N, const vector<ll> weight)
      : num_components(N), parent(std::vector<int>(N)), weight(weight),
        rank(std::vector<int>(N, 0)) {
    for (int i = 0; i < N; i++)
      parent[i] = i;
  }

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

  int size(int x) { return weight[find(x)]; }

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

  void unite(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return;

    num_components--;
    if (rank[x] < rank[y]) {
      weight[y] += weight[x];
      parent[x] = y;
    } else {
      weight[x] += weight[y];
      parent[y] = x;
      if (rank[x] == rank[y])
        rank[y]++;
    }
  }

  int count() { return num_components; }
};

int main(int argc, char *argv[]) {
  int N, M;
  cin >> N >> M;
  vector<ll> WN(N);
  REP(i, N) cin >> WN[i];
  vector<pair<ll, pair<int, int>>> E(M);
  REP(i, M) {
    int a, b, y;
    cin >> a >> b >> y;
    a--;
    b--;
    E[i] = make_pair(y, make_pair(a, b));
  }
  sort(ALL(E));
  UnionFind uf(N, WN);

  vector<tuple<ll, ll, ll>> op;
  vector<int> valid(N - 1, true);
  vector<int> last_op(N, -1);
  vector<vector<int>> ch_ops(N - 1);

  if (N == 1) {
    cout << 0 << endl;
    return 0;
  }

  REP(i, M) {
    int w = E[i].first;
    // cout << w << endl;
    int a = E[i].second.first;
    int b = E[i].second.second;
    if (uf.same(a, b)) {
      continue;
    }
    // cout << "merge: " << a << " " << b << endl;
    a = uf.find(a);
    b = uf.find(b);
    int c = op.size();
    if (last_op[a] != -1) {
      ch_ops[c].push_back(last_op[a]);
    }
    if (last_op[b] != -1) {
      ch_ops[c].push_back(last_op[b]);
    }
    op.push_back(make_tuple(uf.weight[a], uf.weight[b], i));
    uf.unite(a, b);
    last_op[uf.find(a)] = c;
  }
  assert(int(op.size()) == N - 1);
  queue<int> que;
  que.push(N - 2);
  while (!que.empty()) {
    int i = que.front();
    que.pop();
    ll s0 = get<0>(op[i]);
    ll s1 = get<1>(op[i]);
    ll w = E[get<2>(op[i])].first;
    // cout << i << " " << w << " " << s0 << " "<< s1 << endl;
    if (w > s0 + s1) {
      valid[i] = false;
      for (int j : ch_ops[i]) {
        que.push(j);
      }
    }
  }
  UnionFind new_uf = UnionFind(N, WN);
  REP(i, N - 1) {
    if (valid[i]) {
      int eid = get<2>(op[i]);
      int a = E[eid].second.first;
      int b = E[eid].second.second;
      new_uf.unite(a, b);
    }
  }
  ll res = 0;
  REP(i, M) {
    int w = E[i].first;
    int a = new_uf.find(E[i].second.first);
    int b = new_uf.find(E[i].second.second);
    if (!new_uf.same(a, b) || w > new_uf.weight[a]) {
      res++;
    }
  }
  cout << res << endl;
  return 0;
}