Problem p02364 (C++) — buggy submission, verdict: Runtime Error

#include <bits/stdc++.h>

using namespace std;

struct UnionFind {
  vector<int> data;

  UnionFind(int sz) { data.assign(sz, -1); }

  bool unite(int x, int y) {
    x = find(x), y = find(y);
    if (x == y)
      return (false);
    if (data[x] > data[y])
      swap(x, y);
    data[x] += data[y];
    data[y] = x;
    return (true);
  }

  int find(int k) {
    if (data[k] < 0)
      return (k);
    return (data[k] = find(data[k]));
  }

  int size(int k) { return (-data[find(k)]); }
};

template <typename T>
T boruvka(int N, function<vector<pair<T, int>>(int, vector<int>)> f) {
  vector<int> rev(N), belong(N);
  UnionFind uf(N);
  T ret = T();
  while (uf.size(0) != N) {
    int ptr = 0;
    for (int i = 0; i < N; i++) {
      if (uf.find(i) == i) {
        belong[i] = ptr++;
        rev[belong[i]] = i;
      }
    }
    for (int i = 0; i < N; i++) {
      belong[i] = belong[uf.find(i)];
    }
    auto v = f(ptr, belong);
    assert(v.size() == ptr);
    bool update = false;
    for (int i = 0; i < ptr; i++) {
      if (~v[i].second && uf.unite(rev[i], rev[v[i].second])) {
        ret += v[i].first;
        update = true;
      }
    }
    if (!update)
      return -1; // notice!!
  }
  return ret;
}

template <typename T> struct edge {
  int src, to;
  T cost;

  edge(int to, T cost) : src(-1), to(to), cost(cost) {}

  edge(int src, int to, T cost) : src(src), to(to), cost(cost) {}

  edge &operator=(const int &x) {
    to = x;
    return *this;
  }

  operator int() const { return to; }
};

template <typename T> using Edges = vector<edge<T>>;
template <typename T> using WeightedGraph = vector<Edges<T>>;
using UnWeightedGraph = vector<vector<int>>;
template <typename T> using Matrix = vector<vector<T>>;

int main() {
  int V, E;
  cin >> V >> E;
  Edges<int> g;
  for (int i = 0; i < E; i++) {
    int x, y, z;
    cin >> x >> y >> z;
    g.emplace_back(x, y, z);
  }
  const int INF = 1 << 30;
  function<vector<pair<int, int>>(int, vector<int>)> f =
      [&](int sz, vector<int> belong) {
        vector<pair<int, int>> ret(sz, {INF, -1});
        for (auto &e : g) {
          if (belong[e.src] == belong[e.to])
            continue;
          ret[belong[e.src]] =
              min(ret[belong[e.src]], make_pair(e.cost, belong[e.to]));
          ret[belong[e.to]] =
              min(ret[belong[e.to]], make_pair(e.cost, belong[e.src]));
        }
        return ret;
      };
  cout << boruvka(V, f) << endl;
}


Error: -6
ea4a948a-2097-484e-ad0c-8c59c5178ec7.out: /home/alex/Documents/bug-detection/input/Project_CodeNet/data/p02364/C++/s111802414.cpp:47: T boruvka(int, std::function<std::vector<std::pair<_BiIter, int> >(int, std::vector<int>)>) [with T = int]: Assertion `v.size() == belong.size()' failed.


Fixed submission:

#include <bits/stdc++.h>

using namespace std;

struct UnionFind {
  vector<int> data;

  UnionFind(int sz) { data.assign(sz, -1); }

  bool unite(int x, int y) {
    x = find(x), y = find(y);
    if (x == y)
      return (false);
    if (data[x] > data[y])
      swap(x, y);
    data[x] += data[y];
    data[y] = x;
    return (true);
  }

  int find(int k) {
    if (data[k] < 0)
      return (k);
    return (data[k] = find(data[k]));
  }

  int size(int k) { return (-data[find(k)]); }
};

template <typename T>
T boruvka(int N, function<vector<pair<T, int>>(int, vector<int>)> f) {
  vector<int> rev(N), belong(N);
  UnionFind uf(N);
  T ret = T();
  while (uf.size(0) != N) {
    int ptr = 0;
    for (int i = 0; i < N; i++) {
      if (uf.find(i) == i) {
        belong[i] = ptr++;
        rev[belong[i]] = i;
      }
    }
    for (int i = 0; i < N; i++)
      belong[i] = belong[uf.find(i)];
    auto v = f(ptr, belong);
    assert(v.size() == belong.size());
    bool update = false;
    for (int i = 0; i < belong.size(); i++) {
      if (~rev[v[i].second] && uf.unite(rev[i], rev[v[i].second])) {
        ret += v[i].first;
        update = true;
      }
    }
    if (!update)
      return -1; // notice!!
  }
  return ret;
}

template <typename T> struct edge {
  int src, to;
  T cost;

  edge(int to, T cost) : src(-1), to(to), cost(cost) {}

  edge(int src, int to, T cost) : src(src), to(to), cost(cost) {}

  edge &operator=(const int &x) {
    to = x;
    return *this;
  }

  operator int() const { return to; }
};

template <typename T> using Edges = vector<edge<T>>;
template <typename T> using WeightedGraph = vector<Edges<T>>;
using UnWeightedGraph = vector<vector<int>>;
template <typename T> using Matrix = vector<vector<T>>;

int main() {
  int V, E;
  cin >> V >> E;
  Edges<int> g;
  for (int i = 0; i < E; i++) {
    int x, y, z;
    cin >> x >> y >> z;
    g.emplace_back(x, y, z);
  }
  const int INF = 1 << 30;
  function<vector<pair<int, int>>(int, vector<int>)> f =
      [&](int sz, vector<int> belong) {
        vector<pair<int, int>> ret(sz, {INF, -1});
        for (auto &e : g) {
          ret[belong[e.src]] =
              min(ret[belong[e.src]], make_pair(e.cost, belong[e.to]));
          ret[belong[e.to]] =
              min(ret[belong[e.to]], make_pair(e.cost, belong[e.src]));
        }
        return ret;
      };
  cout << boruvka(V, f) << endl;
}