Problem p03143 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
#define int long long
/*
計算量気にしなければ貪欲に消せばいい
はみ出しているのがあったら消去
非連結になったらはみ出し確認
明らかに計算量が死ぬ
*/
// BEGIN CUT HERE
struct UnionFind {
  int num;             // 連結成分の数
  vector<int> r, p, w; // そのグループのサイズ,自分の親っぽいやつ
  UnionFind() {}
  UnionFind(int n) : num(n), r(n, 1), w(n), p(n, 0) {
    iota(p.begin(), p.end(), 0);
  }

  int find(int x) { // どのグループに所属するか
    return (x == p[x]
                ? x
                : p[x] =
                      find(p[x])); // xがグループの名前と一致するまでxを親にする
  }

  void init(vector<int> v) { w = v; }

  bool same(int x, int y) { // 同じグループかどうか
    return find(x) == find(y);
  }

  void unite(int x, int y) { // xとyを同じグループにする
    x = find(x);
    y = find(y); // xとyのグループの名前をどっちかが変える
    if (x == y)
      return;
    if (r[x] < r[y])
      swap(x, y); // サイズが大きい方をxとする
    r[x] += r[y]; // yの親をxにする（今までyだったグループ名がxになる）
    p[y] = x;
    w[x] += w[y];
    num--;
  }

  int size(int x) { // グループの大きさ
    return r[find(x)];
  }

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

  int count() const { // グループの数
    return num;
  }
};

typedef pair<int, int> P;
typedef pair<P, P> PP;

signed main() {
  int N, M;
  cin >> N >> M;
  vector<int> w(N);
  for (int i = 0; i < N; i++)
    cin >> w[i];
  vector<PP> e(M);
  vector<P> edge[N];
  for (int i = 0; i < M; i++) {
    int a, b, c;
    cin >> a >> b >> c;
    a--;
    b--;
    e[i] = PP(P(c, 0), P(a, b));
    edge[a].push_back(P(b, c));
    edge[b].push_back(P(a, c));
  }
  sort(e.begin(), e.end());
  UnionFind uf(N);
  uf.init(w);
  for (int i = 0; i < N; i++) {
    int W = e[i].first.first, U = e[i].second.first, V = e[i].second.second;
    uf.unite(U, V);
    if (W <= uf.weight(U))
      e[i].first.second = 1;
  }
  sort(e.rbegin(), e.rend());
  set<P> s;
  for (int i = 0; i < M; i++) {
    int W = e[i].first.first, A = e[i].first.second, U = e[i].second.first,
        V = e[i].second.second;
    if (s.count(P(U, V)))
      continue;
    if (A == 0)
      continue;
    s.count(P(U, V));
    queue<int> que;
    que.push(U);
    que.push(V);
    while (que.size()) {
      int p = que.front();
      que.pop();
      for (auto q : edge[p]) {
        int t = q.first, WE = q.second;
        if (s.count(P(min(p, t), max(p, t))))
          continue;
        if (WE > W)
          continue;
        s.insert(P(min(p, t), max(p, t)));
        que.push(t);
      }
    }
  }
  cout << M - s.size() << endl;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
#define int long long
/*
計算量気にしなければ貪欲に消せばいい
はみ出しているのがあったら消去
非連結になったらはみ出し確認
明らかに計算量が死ぬ
*/
// BEGIN CUT HERE
struct UnionFind {
  int num;             // 連結成分の数
  vector<int> r, p, w; // そのグループのサイズ,自分の親っぽいやつ
  UnionFind() {}
  UnionFind(int n) : num(n), r(n, 1), w(n), p(n, 0) {
    iota(p.begin(), p.end(), 0);
  }

  int find(int x) { // どのグループに所属するか
    return (x == p[x]
                ? x
                : p[x] =
                      find(p[x])); // xがグループの名前と一致するまでxを親にする
  }

  void init(vector<int> v) { w = v; }

  bool same(int x, int y) { // 同じグループかどうか
    return find(x) == find(y);
  }

  void unite(int x, int y) { // xとyを同じグループにする
    x = find(x);
    y = find(y); // xとyのグループの名前をどっちかが変える
    if (x == y)
      return;
    if (r[x] < r[y])
      swap(x, y); // サイズが大きい方をxとする
    r[x] += r[y]; // yの親をxにする（今までyだったグループ名がxになる）
    p[y] = x;
    w[x] += w[y];
    num--;
  }

  int size(int x) { // グループの大きさ
    return r[find(x)];
  }

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

  int count() const { // グループの数
    return num;
  }
};

typedef pair<int, int> P;
typedef pair<P, P> PP;

signed main() {
  int N, M;
  cin >> N >> M;
  vector<int> w(N);
  for (int i = 0; i < N; i++)
    cin >> w[i];
  vector<PP> e(M);
  vector<P> edge[N];
  for (int i = 0; i < M; i++) {
    int a, b, c;
    cin >> a >> b >> c;
    a--;
    b--;
    e[i] = PP(P(c, 0), P(a, b));
    edge[a].push_back(P(b, c));
    edge[b].push_back(P(a, c));
  }
  sort(e.begin(), e.end());
  UnionFind uf(N);
  uf.init(w);
  for (int i = 0; i < M; i++) {
    int W = e[i].first.first, U = e[i].second.first, V = e[i].second.second;
    uf.unite(U, V);
    if (W <= uf.weight(U))
      e[i].first.second = 1;
  }
  sort(e.rbegin(), e.rend());
  set<P> s;
  for (int i = 0; i < M; i++) {
    int W = e[i].first.first, A = e[i].first.second, U = e[i].second.first,
        V = e[i].second.second;
    if (s.count(P(U, V)))
      continue;
    if (A == 0)
      continue;
    s.count(P(U, V));
    queue<int> que;
    que.push(U);
    que.push(V);
    while (que.size()) {
      int p = que.front();
      que.pop();
      for (auto q : edge[p]) {
        int t = q.first, WE = q.second;
        if (s.count(P(min(p, t), max(p, t))))
          continue;
        if (WE > W)
          continue;
        s.insert(P(min(p, t), max(p, t)));
        que.push(t);
      }
    }
  }
  cout << M - s.size() << endl;
}