Problem p03143 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;

#define int long long
#define double long double
#define rep(i, n) for (int i = 0; i < (n); i++)
#define reps(i, n) for (int i = 1; i <= (n); i++)
#define all(x) (x).begin(), (x).end()
#define uniq(x) (x).erase(unique(all(x)), (x).end())
#define bit(n) (1LL << (n))
#define dump(x) cerr << #x " = " << (x) << endl
using vint = vector<int>;
using vvint = vector<vint>;
using pint = pair<int, int>;
using vpint = vector<pint>;
template <typename T>
using priority_queue_rev = priority_queue<T, vector<T>, greater<T>>;
constexpr double pi = 3.1415926535897932384626433832795028;
constexpr int dy[8] = {0, 1, 0, -1, 1, 1, -1, -1};
constexpr int dx[8] = {1, 0, -1, 0, 1, -1, -1, 1};
int gcd(int a, int b) {
  while (b) {
    swap(a %= b, b);
  }
  return a;
}
int lcm(int a, int b) { return a / gcd(a, b) * b; }
template <typename T1, typename T2> bool chmax(T1 &a, const T2 &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T1, typename T2> bool chmin(T1 &a, const T2 &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}
template <typename T1, typename T2>
ostream &operator<<(ostream &os, const pair<T1, T2> &rhs) {
  os << "(" << rhs.first << ", " << rhs.second << ")";
  return os;
}
template <typename T> ostream &operator<<(ostream &os, const vector<T> &rhs) {
  os << "{";
  for (auto itr = rhs.begin(); itr != rhs.end(); itr++) {
    os << *itr << (next(itr) != rhs.end() ? ", " : "");
  }
  os << "}";
  return os;
}
struct setup {
  static constexpr int PREC = 20;
  setup() {
    cout << fixed << setprecision(PREC);
    cerr << fixed << setprecision(PREC);
  };
} setup;

struct union_find_tree {
  int n;
  std::vector<int> par, size;
  vint last_edge, weight_sum;
  union_find_tree(int n) : n(n), par(n), size(n), last_edge(n), weight_sum(n) {
    for (int i = 0; i < n; i++) {
      par[i] = i;
      size[i] = 1;
      last_edge[i] = -1;
    }
  }
  int get_root(int i) {
    if (par[i] == i) {
      return i;
    } else {
      return par[i] = get_root(par[i]);
    }
  }
  void unite(int i, int j, int edge) {
    if ((i = get_root(i)) == (j = get_root(j))) {
      return;
    }
    if (size[i] < size[j]) {
      std::swap(i, j);
    }
    par[j] = i;
    last_edge[i] = edge;
    weight_sum[i] += weight_sum[j];
    size[i] += size[j];
  }
  bool are_same(int i, int j) { return get_root(i) == get_root(j); }
  int get_size(int i) { return size[get_root(i)]; }
};

int N, M;
int X[110000];
int A[110000], B[110000], Y[110000];
pair<int, pint> edges[110000];
vint edge_tree[110000];
bool prev_ok[110000], final_ok[110000];

void dfs(int idx) {
  if (idx == -1) {
    return;
  }
  if (prev_ok[idx]) {
    return;
  }
  final_ok[idx] = false;
  rep(i, edge_tree[idx].size()) { dfs(edge_tree[idx][i]); }
}

signed main() {
  cin >> N >> M;
  reps(i, N) { cin >> X[i]; }
  rep(i, M) {
    cin >> A[i] >> B[i] >> Y[i];
    edges[i] = {Y[i], {A[i], B[i]}};
  }
  union_find_tree uft0(N + 1);
  reps(i, N) { uft0.weight_sum[i] = X[i]; }
  sort(edges, edges + M);
  int le;
  rep(i, M) {
    int a = edges[i].second.first, b = edges[i].second.second,
        y = edges[i].first;
    a = uft0.get_root(a), b = uft0.get_root(b);
    if (uft0.are_same(a, b)) {
      continue;
    }
    le = i;
    final_ok[i] = true;
    edge_tree[i].push_back(uft0.last_edge[a]);
    edge_tree[i].push_back(uft0.last_edge[b]);
    prev_ok[i] = y <= uft0.weight_sum[a] + uft0.weight_sum[b];
    uft0.unite(a, b, i);
  }
  dfs(le);
  union_find_tree uft1(N + 1);
  reps(i, N) { uft1.weight_sum[i] = X[i]; }
  rep(i, M) {
    int a = edges[i].second.first, b = edges[i].second.second;
    if (final_ok[i]) {
      uft1.unite(a, b, i);
    }
  }
  int ans = M;
  rep(i, M) {
    int a = edges[i].second.first, y = edges[i].first;
    a = uft1.get_root(a);
    ans -= y <= uft1.weight_sum[a];
  }
  cout << ans << endl;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;

#define int long long
#define double long double
#define rep(i, n) for (int i = 0; i < (n); i++)
#define reps(i, n) for (int i = 1; i <= (n); i++)
#define all(x) (x).begin(), (x).end()
#define uniq(x) (x).erase(unique(all(x)), (x).end())
#define bit(n) (1LL << (n))
#define dump(x) cerr << #x " = " << (x) << endl
using vint = vector<int>;
using vvint = vector<vint>;
using pint = pair<int, int>;
using vpint = vector<pint>;
template <typename T>
using priority_queue_rev = priority_queue<T, vector<T>, greater<T>>;
constexpr double pi = 3.1415926535897932384626433832795028;
constexpr int dy[8] = {0, 1, 0, -1, 1, 1, -1, -1};
constexpr int dx[8] = {1, 0, -1, 0, 1, -1, -1, 1};
int gcd(int a, int b) {
  while (b) {
    swap(a %= b, b);
  }
  return a;
}
int lcm(int a, int b) { return a / gcd(a, b) * b; }
template <typename T1, typename T2> bool chmax(T1 &a, const T2 &b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T1, typename T2> bool chmin(T1 &a, const T2 &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}
template <typename T1, typename T2>
ostream &operator<<(ostream &os, const pair<T1, T2> &rhs) {
  os << "(" << rhs.first << ", " << rhs.second << ")";
  return os;
}
template <typename T> ostream &operator<<(ostream &os, const vector<T> &rhs) {
  os << "{";
  for (auto itr = rhs.begin(); itr != rhs.end(); itr++) {
    os << *itr << (next(itr) != rhs.end() ? ", " : "");
  }
  os << "}";
  return os;
}
struct setup {
  static constexpr int PREC = 20;
  setup() {
    cout << fixed << setprecision(PREC);
    cerr << fixed << setprecision(PREC);
  };
} setup;

struct union_find_tree {
  int n;
  std::vector<int> par, size;
  vint last_edge, weight_sum;
  union_find_tree(int n) : n(n), par(n), size(n), last_edge(n), weight_sum(n) {
    for (int i = 0; i < n; i++) {
      par[i] = i;
      size[i] = 1;
      last_edge[i] = -1;
    }
  }
  int get_root(int i) {
    if (par[i] == i) {
      return i;
    } else {
      return par[i] = get_root(par[i]);
    }
  }
  void unite(int i, int j, int edge) {
    if ((i = get_root(i)) == (j = get_root(j))) {
      return;
    }
    if (size[i] < size[j]) {
      std::swap(i, j);
    }
    par[j] = i;
    last_edge[i] = edge;
    weight_sum[i] += weight_sum[j];
    size[i] += size[j];
  }
  bool are_same(int i, int j) { return get_root(i) == get_root(j); }
  int get_size(int i) { return size[get_root(i)]; }
};

int N, M;
int X[110000];
int A[110000], B[110000], Y[110000];
pair<int, pint> edges[110000];
vint edge_tree[110000];
bool prev_ok[110000], final_ok[110000];

void dfs(int idx) {
  if (idx == -1) {
    return;
  }
  if (prev_ok[idx]) {
    return;
  }
  final_ok[idx] = false;
  rep(i, edge_tree[idx].size()) { dfs(edge_tree[idx][i]); }
}

signed main() {
  cin >> N >> M;
  if (N == 1) {
    cout << 0 << endl;
    return 0;
  }
  reps(i, N) { cin >> X[i]; }
  rep(i, M) {
    cin >> A[i] >> B[i] >> Y[i];
    edges[i] = {Y[i], {A[i], B[i]}};
  }
  union_find_tree uft0(N + 1);
  reps(i, N) { uft0.weight_sum[i] = X[i]; }
  sort(edges, edges + M);
  int le;
  rep(i, M) {
    int a = edges[i].second.first, b = edges[i].second.second,
        y = edges[i].first;
    a = uft0.get_root(a), b = uft0.get_root(b);
    if (uft0.are_same(a, b)) {
      continue;
    }
    le = i;
    final_ok[i] = true;
    edge_tree[i].push_back(uft0.last_edge[a]);
    edge_tree[i].push_back(uft0.last_edge[b]);
    prev_ok[i] = y <= uft0.weight_sum[a] + uft0.weight_sum[b];
    uft0.unite(a, b, i);
  }
  dfs(le);
  union_find_tree uft1(N + 1);
  reps(i, N) { uft1.weight_sum[i] = X[i]; }
  rep(i, M) {
    int a = edges[i].second.first, b = edges[i].second.second;
    if (final_ok[i]) {
      uft1.unite(a, b, i);
    }
  }
  int ans = M;
  rep(i, M) {
    int a = edges[i].second.first, y = edges[i].first;
    a = uft1.get_root(a);
    ans -= y <= uft1.weight_sum[a];
  }
  cout << ans << endl;
}