Problem p03143 (C++)

Buggy submission (status: Time Limit Exceeded):
/*input
6 10
4 4 1 1 1 7
3 5 19
2 5 20
4 5 8
1 6 16
2 3 9
3 6 16
3 4 1
2 6 20
2 4 19
1 2 9
*/
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
using namespace std;
using namespace __gnu_pbds;
#define REP(i, j, k) for (int i = j; i < k; ++i)
#define RREP(i, j, k) for (int i = j; i >= k; --i)
#define F first
#define S second
#define mp make_pair
#define pb emplace_back
#define PII pair<int, int>
#define DBGG(i, j) cout << i << " " << j << endl
#define DB4(i, j, k, l) cout << i << " " << j << " " << k << " " << l << endl
#define IOS cin.tie(0), cout.sync_with_stdio(0)
#define endl "\n"
///------------------------------------------------------------
#define int long long
#define MAX 100009
#define INF 0x3f3f3f3f

struct djs {
  int x[MAX], v[MAX], siz[];
  void init() { REP(i, 0, MAX) x[i] = i; }
  int Find(int now) { return x[now] == now ? now : Find(x[now]); }
  void Union(int a, int b) {
    a = Find(a), b = Find(b);
    if (v[a] < v[b])
      swap(a, b);
    x[a] = b;
    v[b] += v[a];
  }
  int operator[](int now) { return Find(now); }
} ds;

int n, m, e[MAX][3], x[MAX], in[MAX], s[MAX];
int cut[MAX][2], rrr[MAX][2], is[MAX];
int cmp(int a, int b) { return e[a][2] < e[b][2]; }
int32_t main() {
  IOS;
  cin >> n >> m;
  REP(i, 1, n + 1) cin >> x[i];
  REP(i, 1, m + 1) REP(j, 0, 3) cin >> e[i][j];
  REP(i, 1, m + 1) s[i] = i;
  sort(s + 1, s + 1 + m, cmp);

  ds.init();
  REP(i, 1, n + 1) ds.v[i] = x[i];

  REP(idx, 1, m + 1) {
    int i = s[idx];
    int v1 = e[i][0], v2 = e[i][1];

    if (ds[v1] != ds[v2]) {
      is[i] = 1;
      rrr[i][0] = ds[v1];
      cut[i][0] = ds.v[ds[v1]];
      rrr[i][1] = ds[v2];
      cut[i][1] = ds.v[ds[v2]];
      ds.Union(v1, v2);
    }
    // cout << "\t\t\t"; REP(i , 1 , n + 1) cout << ds.x[i] << " "; cout <<
    // endl; cout << "\t\t\t"; REP(i , 1 , n + 1) cout << ds.v[i] << " "; cout
    // << endl;
  }
  int ans = 0;
  RREP(idx, m, 1) {
    int i = s[idx];
    int v1 = e[i][0], v2 = e[i][1];
    // DB4("edge " , e[i][2] , "vs = " , ds.v[ds[v1]]);
    //         DB4(rrr[i][0] , cut[i][0] , rrr[i][1] , cut[i][1]);
    if (ds[v1] != ds[v2])
      ans++; // , DBGG("HERE" , "");
    else if (ds.v[ds[v1]] < e[i][2]) {
      ans++;
      if (is[i]) {
        ds.x[rrr[i][0]] = rrr[i][0], ds.v[rrr[i][0]] = cut[i][0];
        ds.x[rrr[i][1]] = rrr[i][1], ds.v[rrr[i][1]] = cut[i][1];
        // cout << "\t\t\t"; REP(i , 1 , n + 1) cout << ds.x[i] << " "; cout <<
        // endl;
      }
    }
  }
  cout << ans << endl;
  return 0;
}

Fixed submission:
/*input
6 10
4 4 1 1 1 7
3 5 19
2 5 20
4 5 8
1 6 16
2 3 9
3 6 16
3 4 1
2 6 20
2 4 19
1 2 9
*/
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
using namespace std;
using namespace __gnu_pbds;
#define REP(i, j, k) for (int i = j; i < k; ++i)
#define RREP(i, j, k) for (int i = j; i >= k; --i)
#define F first
#define S second
#define mp make_pair
#define pb emplace_back
#define PII pair<int, int>
#define DBGG(i, j) cout << i << " " << j << endl
#define DB4(i, j, k, l) cout << i << " " << j << " " << k << " " << l << endl
#define IOS cin.tie(0), cout.sync_with_stdio(0)
#define endl "\n"
///------------------------------------------------------------
#define int long long
#define MAX 100009
#define INF 0x3f3f3f3f

struct djs {
  int x[MAX], v[MAX], siz[];
  void init() { REP(i, 0, MAX) x[i] = i; }
  int Find(int now) { return x[now] == now ? now : Find(x[now]); }
  void Union(int a, int b) {
    a = Find(a), b = Find(b);
    if (v[a] > v[b])
      swap(a, b);
    x[a] = b;
    v[b] += v[a];
  }
  int operator[](int now) { return Find(now); }
} ds;

int n, m, e[MAX][3], x[MAX], in[MAX], s[MAX];
int cut[MAX][2], rrr[MAX][2], is[MAX];
int cmp(int a, int b) { return e[a][2] < e[b][2]; }
int32_t main() {
  IOS;
  cin >> n >> m;
  REP(i, 1, n + 1) cin >> x[i];
  REP(i, 1, m + 1) REP(j, 0, 3) cin >> e[i][j];
  REP(i, 1, m + 1) s[i] = i;
  sort(s + 1, s + 1 + m, cmp);

  ds.init();
  REP(i, 1, n + 1) ds.v[i] = x[i];

  REP(idx, 1, m + 1) {
    int i = s[idx];
    int v1 = e[i][0], v2 = e[i][1];

    if (ds[v1] != ds[v2]) {
      is[i] = 1;
      rrr[i][0] = ds[v1];
      cut[i][0] = ds.v[ds[v1]];
      rrr[i][1] = ds[v2];
      cut[i][1] = ds.v[ds[v2]];
      ds.Union(v1, v2);
    }
    // cout << "\t\t\t"; REP(i , 1 , n + 1) cout << ds.x[i] << " "; cout <<
    // endl; cout << "\t\t\t"; REP(i , 1 , n + 1) cout << ds.v[i] << " "; cout
    // << endl;
  }
  int ans = 0;
  RREP(idx, m, 1) {
    int i = s[idx];
    int v1 = e[i][0], v2 = e[i][1];
    // DB4("edge " , e[i][2] , "vs = " , ds.v[ds[v1]]);
    //         DB4(rrr[i][0] , cut[i][0] , rrr[i][1] , cut[i][1]);
    if (ds[v1] != ds[v2])
      ans++; // , DBGG("HERE" , "");
    else if (ds.v[ds[v1]] < e[i][2]) {
      ans++;
      if (is[i]) {
        ds.x[rrr[i][0]] = rrr[i][0], ds.v[rrr[i][0]] = cut[i][0];
        ds.x[rrr[i][1]] = rrr[i][1], ds.v[rrr[i][1]] = cut[i][1];
        // cout << "\t\t\t"; REP(i , 1 , n + 1) cout << ds.x[i] << " "; cout <<
        // endl;
      }
    }
  }
  cout << ans << endl;
  return 0;
}