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

#include <bits/stdc++.h>
using Int = int64_t;
using UInt = uint64_t;
using C = std::complex<double>;
#define rep(i, n) for (Int i = 0; i < (Int)(n); ++i)
#define guard(x)                                                               \\
  if (not(x))                                                                  \\
    continue;
#ifndef LOCAL_
#define fprintf                                                                \\
  if (false)                                                                   \\
  fprintf
#endif

template <typename T>
using RQ = std::priority_queue<T, std::vector<T>, std::greater<T>>;

struct Tree {
  Tree(Int n) : ps(n, -1) {}
  std::vector<Int> ps;
  Int root(Int x) {
    if (ps[x] == -1)
      return x;
    return ps[x] = root(ps[x]);
  }
  void connect(Int parent, Int child) {
    parent = root(parent), child = root(child);
    if (parent == child)
      return;
    ps[root(child)] = root(parent);
  }
};

int main() {
  Int v, e, r;
  std::cin >> v >> e >> r;
  std::vector<Int> ss(e), ts(e), ws(e);
  rep(i, e) std::cin >> ss[i] >> ts[i] >> ws[i];
  std::vector<RQ<std::pair<Int, Int>>> qs(v);
  std::vector<Int> os(v);
  rep(i, e) {
    Int s = ss[i], t = ts[i], w = ws[i];
    qs[t].emplace(w, s);
  }
  Tree tree(v);
  auto f = [&]() -> Int {
    Int res = 0;
    std::vector<Int> pws(v), pis(v);
    std::vector<bool> finished(v);
    finished[r] = true;
    rep(i, v) {
      Int a = tree.root(i);
      guard(not finished[a]);
      std::vector<Int> ys;
      std::vector<bool> visited(v);
      while (not finished[a]) {
        ys.emplace_back(a);
        visited[a] = true;
        if (qs[a].empty())
          return -1;
        Int w, p;
        do {
          std::tie(w, p) = qs[a].top();
          qs[a].pop();
          w += os[a];
          p = tree.root(p);
        } while (p == tree.root(a));
        pws[a] = w;
        pis[a] = p;
        res += w;
        if (not visited[p]) {
          a = p;
        } else {
          Int nr = a;
          for (Int b = pis[a]; b != a; b = pis[b]) {
            if (qs[nr].size() < qs[b].size())
              nr = b;
          }
          os[nr] -= pws[nr];
          for (Int b = pis[nr]; b != nr; b = pis[b]) {
            tree.connect(nr, b);
            while (not qs[b].empty()) {
              Int w2, p2;
              std::tie(w2, p2) = qs[b].top();
              qs[b].pop();
              Int nw = w2 + os[b] - pws[b] - os[nr];
              qs[nr].emplace(nw, p2);
            }
          }
          a = nr;
        }
      }
      for (Int y : ys)
        finished[y] = true;
    }
    return res;
  };
  printf("%ld\\n", f());
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using Int = int64_t;
using UInt = uint64_t;
using C = std::complex<double>;
#define rep(i, n) for (Int i = 0; i < (Int)(n); ++i)
#define guard(x)                                                               \\
  if (not(x))                                                                  \\
    continue;
#ifndef LOCAL_
#define fprintf                                                                \\
  if (false)                                                                   \\
  fprintf
#endif

template <typename T>
using RQ = std::priority_queue<T, std::vector<T>, std::greater<T>>;

struct Tree {
  Tree(Int n) : ps(n, -1) {}
  std::vector<Int> ps;
  Int root(Int x) {
    if (ps[x] == -1)
      return x;
    return ps[x] = root(ps[x]);
  }
  void connect(Int parent, Int child) {
    parent = root(parent), child = root(child);
    if (parent == child)
      return;
    ps[root(child)] = root(parent);
  }
};

int main() {
  Int v, e, r;
  std::cin >> v >> e >> r;
  std::vector<Int> ss(e), ts(e), ws(e);
  rep(i, e) std::cin >> ss[i] >> ts[i] >> ws[i];
  std::vector<RQ<std::pair<Int, Int>>> qs(v);
  std::vector<Int> os(v);
  rep(i, e) {
    Int s = ss[i], t = ts[i], w = ws[i];
    qs[t].emplace(w, s);
  }
  Tree tree(v);
  auto f = [&]() -> Int {
    Int res = 0;
    std::vector<Int> pws(v), pis(v);
    std::vector<bool> finished(v);
    std::vector<bool> visited(v);
    finished[r] = true;
    rep(i, v) {
      Int a = tree.root(i);
      guard(not finished[a]);
      std::vector<Int> ys;
      while (not finished[a]) {
        ys.emplace_back(a);
        visited[a] = true;
        if (qs[a].empty())
          return -1;
        Int w, p;
        do {
          std::tie(w, p) = qs[a].top();
          qs[a].pop();
          w += os[a];
          p = tree.root(p);
        } while (p == tree.root(a));
        pws[a] = w;
        pis[a] = p;
        res += w;
        if (not visited[p]) {
          a = p;
        } else {
          Int nr = a;
          for (Int b = pis[a]; b != a; b = pis[b]) {
            if (qs[nr].size() < qs[b].size())
              nr = b;
          }
          os[nr] -= pws[nr];
          for (Int b = pis[nr]; b != nr; b = pis[b]) {
            tree.connect(nr, b);
            while (not qs[b].empty()) {
              Int w2, p2;
              std::tie(w2, p2) = qs[b].top();
              qs[b].pop();
              Int nw = w2 + os[b] - pws[b] - os[nr];
              qs[nr].emplace(nw, p2);
            }
          }
          a = nr;
        }
      }
      for (Int y : ys)
        finished[y] = true;
    }
    return res;
  };
  printf("%ld\\n", f());
}