Problem p02365 (C++) — buggy submission, verdict: Time Limit Exceeded

#include <bits/stdc++.h>
using namespace std;
#define MAXN 100000
// Edmonds minimal branching algorithm
typedef struct EdgeStruct Edge;
vector<Edge> incoming[2 * MAXN]; // Stores incoming edges
struct UFVal // Disjoint set union find data structure, used for supernodes
{
  int rep[2 * MAXN]; // Stores the highest node we know on the path from this
                     // node to its root
  int val[2 * MAXN]; // Stores sum of all values on the path from the node to
                     // its root, excluding the value of the root
  int valatroot[2 * MAXN]; // Stores the value at a root, or what value it had
                           // when it was a root
  int upto;
  UFVal() // Initialises the union find for MAXN elements
  {
    upto = MAXN;
    for (int i = 0; i < 2 * MAXN; i++)
      rep[i] = i;
  }
  int findrep(int a) // Finds the root of the tree that a is in. Updates the
                     // value of a as well
  {
    if (rep[a] == a)
      return a;
    int b = rep[a];
    rep[a] = findrep(b);
    val[a] += val[b] + valatroot[a];
    valatroot[a] = 0;
    return rep[a];
  }
  bool connected(int a, int b) // Are a and b connected?
  {
    return findrep(a) == findrep(b);
  }
  void merge(int a, int b) // Joins the sets containing a and b
  {
    // Set a and b to their roots
    a = findrep(a);
    b = findrep(b);
    assert(a != b);
    rep[a] = b;
  }
  void addval(int a, int val) // Adds val to the set containing a
  {
    a = findrep(a);
    valatroot[a] += val;
  }
  int findval(int a) // Returns the value of node a
  {
    int b = findrep(a);
    return val[a] + valatroot[b];
  }
};
struct UF // Simpler union find data structure that does not need to support
          // adding a value. Used to see if a path exists between two nodes
{
  int rep[MAXN];
  UF() {
    for (int i = 0; i < MAXN; i++)
      rep[i] = i;
  }
  int findrep(int a) {
    if (rep[a] == a)
      return a;
    return rep[a] = findrep(rep[a]);
  }
  bool connected(int a, int b) { return findrep(a) == findrep(b); }
  void merge(int a, int b) { rep[findrep(a)] = findrep(b); }
};
UFVal supernodes; // Stores all the compacted 'supernodes'
UF trees;         // Stores all the trees that have been constructed so far
struct EdgeStruct // Stores a directed edge from u to v
{
  int u, v;
  int weight;
  int val() // Returns the corrected weight of this edge, i.e. weight - the
            // value subtracted from all edges going into v
  {
    return weight - supernodes.findval(v);
  }
  bool operator==(const EdgeStruct e) const {
    return e.u == u && v == e.v && weight == e.weight;
  }
};
int n, m;
int ans;            // Stores the weight of the minimum spanning tree
queue<int> roots;   // Stores all the roots to be processed
Edge par[2 * MAXN]; // Stores the parent edge of a node, used to find cycles
int main() {
  // Scan in input
  scanf("%d%d", &n, &m);
  int r;
  scanf("%d", &r);
  for (int i = 0; i < m; i++) {
    int a, b;
    int c;
    scanf("%d%d%d", &a, &b, &c);
    Edge e;
    e.u = a;
    e.v = b;
    e.weight = c;
    incoming[b].push_back(e);
  }
  // Push all nodes onto the queue to be processed, except the root
  for (int i = 0; i < n; i++) {
    if (i != r)
      roots.push(i);
  }
  // Run the algorithm
  while (!roots.empty()) {
    int a = roots.front();
    roots.pop();
    if (a != supernodes.findrep(a))
      continue; // Ensure a is the root
    // Find the minimum incoming edge
    Edge e;
    while (true) {
      if (incoming[a].empty())
        assert(false); // If this is the case, branching is not possible
      e = incoming[a][0];
      for (int i = 1; i < incoming[a].size(); i++) {
        if (incoming[a][i].val() < e.val())
          e = incoming[a][i];
      }
      // Must remove this edge
      for (int i = 0; i < incoming[a].size(); i++) {
        if (incoming[a][i] == e) {
          incoming[a].erase(incoming[a].begin() + i,
                            incoming[a].begin() + i + 1);
          break;
        }
      }
      if (!supernodes.connected(e.u, e.v))
        break; // We have found the edge!
    }
    ans += e.val(); // Add value to the answer
    par[a] = e;
    if (trees.findrep(e.u) != trees.findrep(e.v)) // They are in different trees
    {
      trees.merge(e.u, e.v);
    } else // They are in the same tree, therefore a cycle was formed. Contract
           // into a supernode
    {
      Edge mxedge = e; // Maximum edge on the cycle
      int b = supernodes.findrep(e.u);
      while (b != a) // Find the cycle, up
      {
        if (par[b].val() > mxedge.val())
          mxedge = par[b]; // Update maximum edge if needed
        b = supernodes.findrep(par[b].u);
      }
      // Remove mxedge from cost
      ans -= mxedge.val();

      int s = supernodes.upto++; // New supernode
      // Go over the cycle again, merging the incoming edge lists
      b = supernodes.findrep(e.u);
      while (supernodes.findrep(b) != s) {
        // Subtract weight of par[b] - mxedge from the weight of all edges going
        // into b
        supernodes.addval(b, par[b].val() - mxedge.val());
        // Merge the incoming edges of b into a
        for (auto c : incoming[b]) {
          incoming[s].push_back(c);
        }
        // Merge b and a
        supernodes.merge(b, s);
        b = supernodes.findrep(par[b].u);
      }
      roots.push(s);
    }
  }
  // Print the answer
  printf("%d\\n", ans);
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define MAXN 100000
// Edmonds minimal branching algorithm
typedef struct EdgeStruct Edge;
vector<Edge> incoming[2 * MAXN]; // Stores incoming edges
struct UFVal // Disjoint set union find data structure, used for supernodes
{
  int par[MAXN]; // Stores the parent of this node in the UF tree
  int rep[MAXN]; // Stores the highest node we know on the path from this node
                 // to its root
  int rank[MAXN];
  int val[MAXN]; // Stores sum of all values on the path from the node to its
                 // root, excluding the value of the root
  int valatroot[MAXN]; // Stores the value at a root, or what value it had when
                       // it was a root
  UFVal()              // Initialises the union find for MAXN elements
  {
    for (int i = 0; i < MAXN; i++)
      par[i] = rep[i] = i;
  }
  int findrep(int a) // Finds the root of the tree that a is in. Updates the
                     // value of a as well
  {
    int b = rep[a];
    while (par[b] != b) {
      val[a] += valatroot[b]; // B is no longer a root, so add its value to a
      b = par[b];
    }
    rep[a] =
        b; // Store this so we don't need to go through this process next time
    return a;
  }
  bool connected(int a, int b) // Are a and b connected?
  {
    return findrep(a) == findrep(b);
  }
  void merge(int a, int b) // Joins the sets containing a and b
  {
    // Set a and b to their roots
    a = findrep(a);
    b = findrep(b);
    assert(a != b);
    if (rank[a] > rank[b]) // Ensure rank[a] <= rank[b]
    {
      std::swap(a, b);
      std::swap(incoming[a], incoming[b]);
    }
    par[a] = b;
    if (rank[a] == rank[b])
      rank[b]++; // Increment rank of b if needed
  }
  void addval(int a, int val) // Adds val to the set containing a
  {
    a = findrep(a);
    valatroot[a] += val;
  }
  int findval(int a) // Returns the value of node a
  {
    int b = findrep(a);
    return val[a] + valatroot[b];
  }
};
struct UF // Simpler union find data structure that does not need to support
          // adding a value. Used to see if a path exists between two nodes
{
  int rep[MAXN];
  UF() {
    for (int i = 0; i < MAXN; i++)
      rep[i] = i;
  }
  int findrep(int a) {
    if (rep[a] == a)
      return a;
    return rep[a] = findrep(rep[a]);
  }
  bool connected(int a, int b) { return findrep(a) == findrep(b); }
  void merge(int a, int b) { rep[findrep(a)] = findrep(b); }
};
UFVal supernodes; // Stores all the compacted 'supernodes'
UF trees;         // Stores all the trees that have been constructed so far
struct EdgeStruct // Stores a directed edge from u to v
{
  int u, v;
  int weight;
  int val() // Returns the corrected weight of this edge, i.e. weight - the
            // value subtracted from all edges going into v
  {
    return weight - supernodes.findval(v);
  }
  bool operator==(const EdgeStruct e) const {
    return e.u == u && v == e.v && weight == e.weight;
  }
};
int n, m;
int ans;            // Stores the weight of the minimum spanning tree
queue<int> roots;   // Stores all the roots to be processed
Edge par[2 * MAXN]; // Stores the parent edge of a node, used to find cycles
int main() {
  // Scan in input
  scanf("%d%d", &n, &m);
  int r;
  scanf("%d", &r);
  for (int i = 0; i < m; i++) {
    int a, b;
    int c;
    scanf("%d%d%d", &a, &b, &c);
    Edge e;
    e.u = a;
    e.v = b;
    e.weight = c;
    incoming[b].push_back(e);
  }
  // Push all nodes onto the queue to be processed, except the root
  for (int i = 0; i < n; i++) {
    if (i != r)
      roots.push(i);
  }
  // Run the algorithm
  while (!roots.empty()) {
    int a = roots.front();
    roots.pop();
    a = supernodes.findrep(a); // Ensure a is the root
    // Find the minimum incoming edge
    Edge e;
    while (true) {
      if (incoming[a].empty())
        assert(false); // If this is the case, branching is not possible
      e = incoming[a][0];
      for (int i = 1; i < incoming[a].size(); i++) {
        if (incoming[a][i].val() < e.val())
          e = incoming[a][i];
      }
      if (!supernodes.connected(e.u, e.v))
        break; // We have found the edge!
      // Otherwise, must remove this edge
      for (int i = 0; i < incoming[a].size(); i++) {
        if (incoming[a][i] == e) {
          incoming[a].erase(incoming[a].begin() + i);
          break;
        }
      }
    }
    if (incoming[a].empty())
      assert(false); // If this is the case, branching is not possible
    ans += e.val();  // Add value to the answer
    if (trees.findrep(e.u) != trees.findrep(e.v)) // They are in different trees
    {
      trees.merge(e.u, e.v);
      par[a] = e; // A is now in a tree, update its parent
    } else // They are in the same tree, therefore a cycle was formed. Contract
           // into a supernode
    {
      Edge mxedge = e; // Maximum edge on the cycle
      int b = supernodes.findrep(e.u);
      while (b != a) // Find the cycle, up
      {
        if (par[b].val() > mxedge.val())
          mxedge = par[b]; // Update maximum edge if needed
        b = supernodes.findrep(par[b].u);
      }
      // Go over the cycle again, merging the incoming edge lists
      b = supernodes.findrep(e.u);
      while (b != supernodes.findrep(a)) {
        // Subtract weight of par[b] - mxedge from the weight of all edges going
        // into b
        supernodes.addval(b, par[b].val() - mxedge.val());
        // Merge the incoming edges of b into a
        int s = supernodes.findrep(a);
        for (auto c : incoming[b]) {
          incoming[s].push_back(c);
        }
        // Merge b and a
        supernodes.merge(b, s);
        b = supernodes.findrep(par[b].u);
      }
      roots.push(supernodes.findrep(a));
    }
  }
  // Print the answer
  printf("%d\\n", ans);
}