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

#include <bits/stdc++.h>

using namespace std;

struct Edge {
  int from, to, cost, id;
  char type;
};
bool operator<(const Edge &e1, const Edge &e2) { return e1.cost < e2.cost; }

class UnionFind {
private:
  int sz;
  vector<int> par;

public:
  UnionFind(int n) : sz(n) {
    par = vector<int>(n);
    for (int i = 0; i < n; ++i)
      par[i] = i;
  }
  int root(int a) {
    if (par[a] == a)
      return a;
    return par[a] = root(par[a]);
  }
  void unite(int a, int b) { par[root(a)] = par[root(b)]; }
  bool isUnited(int a, int b) { return root(a) == root(b); }
};

const int inf = 1e9;

void add_edge(vector<vector<Edge>> &G, Edge &e) {
  G[e.from].push_back({e.from, e.to, e.cost, e.id, e.type});
  G[e.to].push_back({e.to, e.from, e.cost, e.id, e.type});
}
void remove_edge(vector<vector<Edge>> &G, Edge &e) {
  int f = 0;
  for (int i = 0; i < G[e.from].size(); ++i) {
    if (G[e.from][i].id == e.id) {
      G[e.from].erase(G[e.from].begin() + i, G[e.from].begin() + i + 1);
      ++f;
      break;
    }
  }
  for (int i = 0; i < G[e.to].size(); ++i) {
    if (G[e.to][i].id == e.id) {
      G[e.to].erase(G[e.to].begin() + i, G[e.to].begin() + i + 1);
      ++f;
      break;
    }
  }
  assert(f == 2);
}
Edge dfs(vector<vector<Edge>> &G, int v, int t, char type, int prev,
         vector<bool> &used) {
  Edge res = {0, 0, -2, -1, 'X'};
  if (used[v])
    assert(false);
  used[v] = true;
  if (v == t) {
    res.cost++;
    return res;
  }
  for (int i = 0; i < G[v].size(); ++i) {
    Edge &e = G[v][i];
    assert(e.from == v);
    if (e.to != prev) {
      Edge r = dfs(G, e.to, t, type, v, used);
      if (r.cost >= -1) {
        res = max(res, r);
        if (e.type == type) {
          res = max(res, e);
        }
      }
    }
  }
  return res;
}
Edge calc(vector<vector<Edge>> &G, Edge &e) {
  vector<bool> used(G.size(), false);
  Edge ee = dfs(G, e.from, e.to, (e.type == 'A' ? 'B' : 'A'), -1, used);
  return ee;
}
int main() {
  int n, m, k;
  while (cin >> n >> m >> k, n || m || k) {
    vector<Edge> ae, be, edges;
    for (int i = 0; i < m; ++i) {
      int u, v, w;
      char l;
      cin >> u >> v >> w >> l;
      --u, --v;
      if (l == 'A') {
        ae.push_back({u, v, w, i, l});
      } else {
        be.push_back({u, v, w, i, l});
      }
      edges.push_back({u, v, w, i, l});
    }
    sort(ae.begin(), ae.end());
    sort(be.begin(), be.end());
    sort(edges.begin(), edges.end());
    vector<vector<Edge>> G(n);
    UnionFind uf(n);
    int ans = 0;
    int cnt = 0, acnt = 0, bcnt = 0;
    set<int> s;
    for (int i = 0; i < edges.size(); ++i) {
      if (cnt >= n - 1)
        break;
      Edge &e = edges[i];
      if (uf.isUnited(e.from, e.to))
        continue;
      uf.unite(e.from, e.to);
      ans += e.cost;
      add_edge(G, e);
      ++cnt;
      s.insert(e.id);
      if (e.type == 'A')
        ++acnt;
      else
        ++bcnt;
    }
    if (cnt < n - 1) {
      cout << -1 << endl;
      continue;
    }
    if (acnt > k) {
      swap(acnt, bcnt);
      swap(ae, be);
      k = n - k - 1;
    }
    // remove use edges
    // int s0 = 0;
    // for (int i = ea.size()-1; i >= 0; --i) {
    //     if (s.count(ea.id) > 0) {
    //         s0 = i+1;
    //         break;
    //     }
    // }
    vector<bool> used(ae.size(), false);
    for (int i = 0; i < ae.size(); ++i) {
      if (s.count(ae[i].id) > 0) {
        used[i] = true;
      }
    }
    assert(k - acnt >= 0);
    for (int t = 0; t < k - acnt; ++t) {
      Edge a = {0, 0, -1, -1, 'X'};
      int b = -1;
      for (int i = 0; i < ae.size(); ++i) {
        if (used[i])
          continue;
        // cout << calc(G, ae[i]) << endl;
        Edge e = calc(G, ae[i]);
        if (e.cost >= 0 &&
            (b < 0 || ae[i].cost - e.cost <= ae[b].cost - a.cost)) {
          a = e;
          b = i;
        }
      }
      if (b < 0) {
        ans = -1;
        break;
      }
      // if (a.cost >= inf) assert(false);
      // cout << a.cost << " " << b << endl;
      used[b] = true;
      remove_edge(G, a);
      add_edge(G, ae[b]);
      ans += ae[b].cost;
      ans -= a.cost;
    }
    cout << ans << endl;
  }
}

Error: -6
a66452cf-1682-4c4d-94fb-8a4eb5482918.out: /home/alex/Documents/bug-detection/input/Project_CodeNet/data/p01090/C++/s173115534.cpp:56: void remove_edge(std::vector<std::vector<Edge> >&, Edge&): Assertion `f == 2' failed.


Fixed submission:

#include <bits/stdc++.h>

using namespace std;

struct Edge {
  int from, to, cost, id;
  char type;
};
bool operator<(const Edge &e1, const Edge &e2) { return e1.cost < e2.cost; }

class UnionFind {
private:
  int sz;
  vector<int> par;

public:
  UnionFind(int n) : sz(n) {
    par = vector<int>(n);
    for (int i = 0; i < n; ++i)
      par[i] = i;
  }
  int root(int a) {
    if (par[a] == a)
      return a;
    return par[a] = root(par[a]);
  }
  void unite(int a, int b) { par[root(a)] = par[root(b)]; }
  bool isUnited(int a, int b) { return root(a) == root(b); }
};

const int inf = 1e9;

void add_edge(vector<vector<Edge>> &G, Edge &e) {
  G[e.from].push_back({e.from, e.to, e.cost, e.id, e.type});
  G[e.to].push_back({e.to, e.from, e.cost, e.id, e.type});
}
void remove_edge(vector<vector<Edge>> &G, Edge &e) {
  int f = 0;
  for (int i = 0; i < G[e.from].size(); ++i) {
    if (G[e.from][i].id == e.id) {
      G[e.from].erase(G[e.from].begin() + i, G[e.from].begin() + i + 1);
      ++f;
      break;
    }
  }
  for (int i = 0; i < G[e.to].size(); ++i) {
    if (G[e.to][i].id == e.id) {
      G[e.to].erase(G[e.to].begin() + i, G[e.to].begin() + i + 1);
      ++f;
      break;
    }
  }
  assert(f == 2);
}
Edge dfs(vector<vector<Edge>> &G, int v, int t, char type, int prev,
         vector<bool> &used) {
  Edge res = {0, 0, -1, -1, 'X'};
  if (used[v])
    assert(false);
  used[v] = true;
  if (v == t) {
    res.cost++;
    return res;
  }
  for (int i = 0; i < G[v].size(); ++i) {
    Edge &e = G[v][i];
    assert(e.from == v);
    if (e.to != prev) {
      Edge r = dfs(G, e.to, t, type, v, used);
      if (r.cost >= 0) {
        res = max(res, r);
        if (e.type == type) {
          res = max(res, e);
        }
      }
    }
  }
  return res;
}
Edge calc(vector<vector<Edge>> &G, Edge &e) {
  vector<bool> used(G.size(), false);
  Edge ee = dfs(G, e.from, e.to, (e.type == 'A' ? 'B' : 'A'), -1, used);
  return ee;
}
int main() {
  int n, m, k;
  while (cin >> n >> m >> k, n || m || k) {
    vector<Edge> ae, be, edges;
    for (int i = 0; i < m; ++i) {
      int u, v, w;
      char l;
      cin >> u >> v >> w >> l;
      --u, --v;
      if (l == 'A') {
        ae.push_back({u, v, w, i, l});
      } else {
        be.push_back({u, v, w, i, l});
      }
      edges.push_back({u, v, w, i, l});
    }
    sort(ae.begin(), ae.end());
    sort(be.begin(), be.end());
    sort(edges.begin(), edges.end());
    vector<vector<Edge>> G(n);
    UnionFind uf(n);
    int ans = 0;
    int cnt = 0, acnt = 0, bcnt = 0;
    set<int> s;
    for (int i = 0; i < edges.size(); ++i) {
      if (cnt >= n - 1)
        break;
      Edge &e = edges[i];
      if (uf.isUnited(e.from, e.to))
        continue;
      uf.unite(e.from, e.to);
      ans += e.cost;
      add_edge(G, e);
      ++cnt;
      s.insert(e.id);
      if (e.type == 'A')
        ++acnt;
      else
        ++bcnt;
    }
    if (cnt < n - 1) {
      cout << -1 << endl;
      continue;
    }
    if (acnt > k) {
      swap(acnt, bcnt);
      swap(ae, be);
      k = n - k - 1;
    }
    // remove use edges
    // int s0 = 0;
    // for (int i = ea.size()-1; i >= 0; --i) {
    //     if (s.count(ea.id) > 0) {
    //         s0 = i+1;
    //         break;
    //     }
    // }
    vector<bool> used(ae.size(), false);
    for (int i = 0; i < ae.size(); ++i) {
      if (s.count(ae[i].id) > 0) {
        used[i] = true;
      }
    }
    assert(k - acnt >= 0);
    for (int t = 0; t < k - acnt; ++t) {
      Edge a = {0, 0, inf - 1, -1, 'X'};
      int b = -1;
      for (int i = 0; i < ae.size(); ++i) {
        if (used[i])
          continue;
        // cout << calc(G, ae[i]) << endl;
        Edge e = calc(G, ae[i]);
        if (e.cost < inf - 1 &&
            (b < 0 || ae[i].cost - e.cost <= ae[b].cost - a.cost)) {
          a = e;
          b = i;
        }
      }
      if (b < 0) {
        ans = -1;
        break;
      }
      // if (a.cost >= inf) assert(false);
      // cout << a.cost << " " << b << endl;
      used[b] = true;
      remove_edge(G, a);
      add_edge(G, ae[b]);
      ans += ae[b].cost;
      ans -= a.cost;
    }
    cout << ans << endl;
  }
}