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

#include <bits/stdc++.h>
using namespace std;
using Int = long long;

template <typename T1, typename T2> inline void chmin(T1 &a, T2 b) {
  if (a > b)
    a = b;
};
template <typename T1, typename T2> inline void chmax(T1 &a, T2 b) {
  if (a < b)
    a = b;
};

template <typename T> struct PrimalDual {
  struct edge {
    Int to;
    T cap, cost;
    Int rev;
    edge() {}
    edge(Int to, T cap, T cost, Int rev)
        : to(to), cap(cap), cost(cost), rev(rev) {}
  };

  T INF;
  vector<vector<edge>> G;
  vector<T> h, dist;
  vector<Int> prevv, preve;

  PrimalDual() {}
  PrimalDual(Int n, T INF)
      : INF(INF), G(n), h(n), dist(n), prevv(n), preve(n) {}

  void add_edge(Int from, Int to, T cap, T cost) {
    G[from].emplace_back(to, cap, cost, G[to].size());
    G[to].emplace_back(from, 0, -cost, G[from].size() - 1);
  }

  T flow(Int s, Int t, T f, Int &ok) {
    T res = 0;
    fill(h.begin(), h.end(), 0);
    while (f > 0) {
      struct P {
        T first;
        Int second;
        P(T first, Int second) : first(first), second(second) {}
        bool operator<(const P &a) const { return first > a.first; }
      };

      priority_queue<P> que;

      fill(dist.begin(), dist.end(), INF);

      dist[s] = 0;
      que.emplace(dist[s], s);
      while (!que.empty()) {
        P p = que.top();
        que.pop();
        Int v = p.second;
        if (dist[v] < p.first)
          continue;
        for (Int i = 0; i < (Int)G[v].size(); i++) {
          edge &e = G[v][i];
          if (e.cap == 0)
            continue;
          if (dist[e.to] > dist[v] + e.cost + h[v] - h[e.to]) {
            dist[e.to] = dist[v] + e.cost + h[v] - h[e.to];
            prevv[e.to] = v;
            preve[e.to] = i;
            que.emplace(dist[e.to], e.to);
          }
        }
      }
      if (dist[t] == INF)
        return ok = 0;

      for (Int v = 0; v < (Int)h.size(); v++)
        h[v] += dist[v];

      T d = f;
      for (Int v = t; v != s; v = prevv[v])
        d = min(d, G[prevv[v]][preve[v]].cap);

      f -= d;
      res += d * h[t];
      for (Int v = t; v != s; v = prevv[v]) {
        edge &e = G[prevv[v]][preve[v]];
        e.cap -= d;
        G[v][e.rev].cap += d;
      }
    }
    ok = 1;
    return res;
  }
};

// INSERT ABOVE HERE
signed main() {
  cin.tie(0);
  ios::sync_with_stdio(0);

  Int n;
  cin >> n;
  const Int INF2 = 1e9;
  const Int INF3 = 1e15;
  PrimalDual<Int> G(n + 1, INF3);
  Int num = 0, ans = 0;
  for (Int i = 0; i < n; i++) {
    Int k;
    cin >> k;
    num += k;
    map<Int, Int> cnt, dst;
    for (Int j = 0; j < k; j++) {
      Int t, c;
      cin >> t >> c;
      t--;
      cnt[t]++;
      ans += c + INF2;
      if (!dst.count(t))
        dst[t] = c;
      chmin(dst[t], c);
    }

    for (auto p : cnt) {
      G.add_edge(i, p.first, cnt[p.first], -INF2);
      G.add_edge(i, p.first, INF2, dst[p.first]);
    }

    G.add_edge(i, n, INF2, 0);
  }

  Int ok = 1;
  ans += G.flow(0, n, INF2, ok);
  cout << ans << endl;
  assert(ok);
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
using Int = long long;

template <typename T1, typename T2> inline void chmin(T1 &a, T2 b) {
  if (a > b)
    a = b;
};
template <typename T1, typename T2> inline void chmax(T1 &a, T2 b) {
  if (a < b)
    a = b;
};

template <typename T> struct PrimalDual {
  struct edge {
    Int to;
    T cap, cost;
    Int rev;
    edge() {}
    edge(Int to, T cap, T cost, Int rev)
        : to(to), cap(cap), cost(cost), rev(rev) {}
  };

  T INF;
  vector<vector<edge>> G;
  vector<T> h, dist;
  vector<Int> prevv, preve;

  PrimalDual() {}
  PrimalDual(Int n, T INF)
      : INF(INF), G(n), h(n), dist(n), prevv(n), preve(n) {}

  void add_edge(Int from, Int to, T cap, T cost) {
    G[from].emplace_back(to, cap, cost, G[to].size());
    G[to].emplace_back(from, 0, -cost, G[from].size() - 1);
  }

  T flow(Int s, Int t, T f, Int &ok) {
    T res = 0;
    fill(h.begin(), h.end(), 0);
    while (f > 0) {
      using P = pair<T, Int>;
      priority_queue<P, vector<P>, greater<P>> que;
      fill(dist.begin(), dist.end(), INF);

      dist[s] = 0;
      que.emplace(dist[s], s);
      while (!que.empty()) {
        P p = que.top();
        que.pop();
        Int v = p.second;
        if (dist[v] < p.first)
          continue;
        for (Int i = 0; i < (Int)G[v].size(); i++) {
          edge &e = G[v][i];
          if (e.cap == 0)
            continue;
          if (dist[e.to] > dist[v] + e.cost + h[v] - h[e.to]) {
            dist[e.to] = dist[v] + e.cost + h[v] - h[e.to];
            prevv[e.to] = v;
            preve[e.to] = i;
            que.emplace(dist[e.to], e.to);
          }
        }
      }
      if (dist[t] == INF)
        return ok = 0;

      for (Int v = 0; v < (Int)h.size(); v++)
        h[v] += dist[v];

      T d = f;
      for (Int v = t; v != s; v = prevv[v])
        d = min(d, G[prevv[v]][preve[v]].cap);

      f -= d;
      res += d * h[t];
      for (Int v = t; v != s; v = prevv[v]) {
        edge &e = G[prevv[v]][preve[v]];
        e.cap -= d;
        G[v][e.rev].cap += d;
      }
    }
    ok = 1;
    return res;
  }
};

// INSERT ABOVE HERE
signed main() {
  cin.tie(0);
  ios::sync_with_stdio(0);

  Int n;
  cin >> n;
  const Int INF2 = 1e9;
  const Int INF3 = 1e15;
  PrimalDual<Int> G(n + 1, INF3);
  Int num = 0, ans = 0;
  for (Int i = 0; i < n; i++) {
    Int k;
    cin >> k;
    num += k;
    map<Int, Int> cnt, dst;
    for (Int j = 0; j < k; j++) {
      Int t, c;
      cin >> t >> c;
      t--;
      cnt[t]++;
      ans += c + INF2;
      if (!dst.count(t))
        dst[t] = c;
      chmin(dst[t], c);
    }

    for (auto p : cnt) {
      G.add_edge(i, p.first, cnt[p.first], -INF2);
      G.add_edge(i, p.first, INF2, dst[p.first]);
    }

    G.add_edge(i, n, INF2, 0);
  }

  Int ok = 1;
  ans += G.flow(0, n, INF2, ok);
  cout << ans << endl;
  assert(ok);
  return 0;
}