Problem p03143 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
using namespace std;
#define int long long
#define rep(i, n) for (int i = 0; i < (n); i++)

class UnionFind {
private:
  vector<int> p, s, w;
  vector<vector<int>> h;

public:
  UnionFind(const vector<int> &v)
      : p(v.size(), -1), s(v.size(), 1), w(v.begin(), v.end()) {}
  void wrap() { h.push_back({}); }
  void unwrap() {
    int k = h.back().size();
    rep(i, k) {
      s[p[h.back()[k - 1 - i]]] -= s[h.back()[k - 1 - i]];
      w[p[h.back()[k - 1 - i]]] -= w[h.back()[k - 1 - i]];
      p[h.back()[k - 1 - i]] = -1;
    }
    h.pop_back();
  }
  int find(int a) {
    if (p[a] < 0) {
      return a;
    }
    return find(p[a]);
  }
  void unite(int a, int b) {
    a = find(a);
    b = find(b);
    if (a != b) {
      if (s[a] >= s[b]) {
        h.back().push_back(b);
        s[a] += s[b];
        w[a] += w[b];
        p[b] = a;
      } else {
        h.back().push_back(a);
        s[b] += s[a];
        w[b] += w[a];
        p[a] = b;
      }
    }
  }
  int get(int a) { return w[find(a)]; }
};

vector<int> v;
vector<pair<int, pair<int, int>>> ed;

vector<int> dfs(int b, int e, UnionFind &uf) {
  if (b == e - 1) {
    uf.wrap();
    uf.unite(ed[b].second.first, ed[b].second.second);
    if (uf.get(ed[b].second.first) >= ed[b].first) {
      uf.unwrap();
      return {b};
    } else {
      uf.unwrap();
      return {};
    }
  } else {
    int m = (b + e) / 2;
    uf.wrap();
    for (int i = b; i < m; i++) {
      uf.unite(ed[i].second.first, ed[i].second.second);
    }
    auto r1 = dfs(m, e, uf);
    uf.unwrap();
    uf.wrap();
    for (auto i : r1) {
      uf.unite(ed[i].second.first, ed[i].second.second);
    }
    auto r2 = dfs(b, m, uf);
    uf.unwrap();
    r2.insert(r2.end(), r1.begin(), r1.end());
    return r2;
  }
}

signed main() {
  int n, m;
  cin >> n >> m;
  rep(i, n) {
    int w;
    cin >> w;
    v.push_back(w);
  }
  rep(i, m) {
    int a, b, y;
    cin >> a >> b >> y;
    ed.push_back(make_pair(y, make_pair(a - 1, b - 1)));
  }
  sort(ed.begin(), ed.end());
  UnionFind uf(v);
  cout << m - dfs(0, m, uf).size() << endl;
}


Fixed submission:
#include <bits/stdc++.h>
using namespace std;
#define int long long
#define rep(i, n) for (int i = 0; i < (n); i++)

class UnionFind {
private:
  vector<int> p, s, w;
  vector<vector<int>> h;

public:
  UnionFind(const vector<int> &v)
      : p(v.size(), -1), s(v.size(), 1), w(v.begin(), v.end()) {}
  void wrap() { h.push_back({}); }
  void unwrap() {
    int k = h.back().size();
    rep(i, k) {
      s[p[h.back()[k - 1 - i]]] -= s[h.back()[k - 1 - i]];
      w[p[h.back()[k - 1 - i]]] -= w[h.back()[k - 1 - i]];
      p[h.back()[k - 1 - i]] = -1;
    }
    h.pop_back();
  }
  int find(int a) {
    if (p[a] < 0) {
      return a;
    }
    return find(p[a]);
  }
  void unite(int a, int b) {
    a = find(a);
    b = find(b);
    if (a != b) {
      if (s[a] >= s[b]) {
        h.back().push_back(b);
        s[a] += s[b];
        w[a] += w[b];
        p[b] = a;
      } else {
        h.back().push_back(a);
        s[b] += s[a];
        w[b] += w[a];
        p[a] = b;
      }
    }
  }
  int get(int a) { return w[find(a)]; }
};

vector<int> v;
vector<pair<int, pair<int, int>>> ed;

vector<int> dfs(int b, int e, UnionFind &uf) {
  if (b == e - 1) {
    uf.wrap();
    uf.unite(ed[b].second.first, ed[b].second.second);
    if (uf.get(ed[b].second.first) >= ed[b].first) {
      uf.unwrap();
      return {b};
    } else {
      uf.unwrap();
      return {};
    }
  } else {
    int m = (b + e) / 2;
    uf.wrap();
    for (int i = b; i < m; i++) {
      uf.unite(ed[i].second.first, ed[i].second.second);
    }
    auto r1 = dfs(m, e, uf);
    uf.unwrap();
    uf.wrap();
    for (auto i : r1) {
      uf.unite(ed[i].second.first, ed[i].second.second);
    }
    auto r2 = dfs(b, m, uf);
    uf.unwrap();
    r2.insert(r2.end(), r1.begin(), r1.end());
    return r2;
  }
}

signed main() {
  int n, m;
  cin >> n >> m;
  rep(i, n) {
    int w;
    cin >> w;
    v.push_back(w);
  }
  rep(i, m) {
    int a, b, y;
    cin >> a >> b >> y;
    ed.push_back(make_pair(y, make_pair(a - 1, b - 1)));
  }
  if (m == 0) {
    cout << 0 << endl;
    return 0;
  }
  sort(ed.begin(), ed.end());
  UnionFind uf(v);
  cout << m - dfs(0, m, uf).size() << endl;
}