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

#include <bits/stdc++.h>

using namespace std;

#define int long long
#define all(v) (v).begin(), (v).end()
#define resz(v, ...) (v).clear(), (v).resize(__VA_ARGS__)
#define reps(i, m, n) for (int i = (int)(m); i < (int)(n); i++)
#define rep(i, n) reps(i, 0, n)

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

using Pi = pair<int, int>;
using Tapris = tuple<int, int, int>;
using vint = vector<int>;

const int inf = 1LL << 55;
const int mod = 1e9 + 7;

struct edge {
  int to, cap, cost, rev;
  edge() {}
  edge(int to, int cap, int cost, int rev)
      : to(to), cap(cap), cost(cost), rev(rev) {}
};

using Graph = vector<vector<edge>>;

struct PrimalDual {
  Graph graph;
  vint mincost, h, prevv, preve;
  PrimalDual(int n) : graph(n), mincost(n), h(n), prevv(n), preve(n) {}
  void add_edge(int from, int to, int cap, int cost) {
    graph[from].emplace_back(to, cap, cost, graph[to].size());
    graph[to].emplace_back(from, 0, -cost, graph[from].size() - 1);
  }
  int mincost_flow(int s, int t, int f) {
    fill(all(h), 0);
    fill(all(prevv), -1);
    fill(all(preve), -1);
    int res = 0;
    while (f > 0) {
      fill(all(mincost), inf);
      mincost[s] = 0;
      priority_queue<Pi, vector<Pi>, greater<Pi>> que;
      que.emplace(0, s);
      while (!que.empty()) {
        int cost, v;
        tie(cost, v) = que.top();
        que.pop();
        if (mincost[v] < cost)
          continue;
        for (int i = 0; i < (int)graph[v].size(); ++i) {
          edge &e = graph[v][i];
          int dual = cost + e.cost + h[v] - h[e.to];
          if (e.cap > 0 && dual < mincost[e.to]) {
            mincost[e.to] = dual;
            prevv[e.to] = v;
            preve[e.to] = i;
            que.emplace(dual, e.to);
          }
        }
      }
      if (mincost[t] == inf)
        return res;
      for (int v = 0; v < (int)graph.size(); ++v)
        h[v] += mincost[v];
      int d = f;
      for (int v = t; v != s; v = prevv[v])
        chmin(d, graph[prevv[v]][preve[v]].cap);
      f -= d;
      res += d * h[t];
      for (int v = t; v != s; v = prevv[v]) {
        edge &e = graph[prevv[v]][preve[v]];
        e.cap -= d;
        graph[e.to][e.rev].cap += d;
      }
    }
    return res;
  }
};

signed main() {
  int n;
  scanf("%lld", &n);
  vector<vint> dp(n, vint(n, inf));
  vint in(n, 0), out(n, 0);
  int sum = 0;
  rep(i, n) dp[i][i] = 0;
  in[0] = inf;
  rep(i, n) {
    int k;
    scanf("%lld", &k);
    rep(j, k) {
      int t, c;
      scanf("%lld %lld", &t, &c);
      --t;
      chmin(dp[i][t], c);
      sum += c;
      ++in[t];
      ++out[i];
    }
  }
  rep(i, n) rep(j, n) rep(k, n) chmin(dp[j][k], dp[j][i] + dp[i][k]);
  vint a, b;
  rep(i, n) {
    if (in[i] > out[i])
      a.push_back(i);
    else if (in[i] < out[i])
      b.push_back(i);
  }
  // cout<<a.size()<<" "<<b.size()<<endl;
  int s = a.size() + b.size(), t = s + 1, sz = t + 1;
  PrimalDual graph(sz);
  rep(i, a.size()) graph.add_edge(s, i, in[a[i]] - out[a[i]], 0);
  int outnum = 0;
  rep(i, b.size()) {
    graph.add_edge(a.size() + i, t, out[b[i]] - in[b[i]], 0);
    outnum += out[b[i]] - in[b[i]];
  }
  int esz = 0;
  /*
  rep(i, a.size()) {
    rep(j, b.size()) {
      if(dp[a[i]][b[j]] < inf) {
        graph.add_edge(i, a.size()+j, inf, dp[a[i]][b[j]]);
        ++esz;
      }
    }
  }
  /*/
  // cout<<outnum<<endl;
  rep(i, b.size()) {
    vector<Pi> dist;
    rep(j, a.size()) {
      if (dp[a[j]][b[i]] < inf)
        dist.emplace_back(dp[a[j]][b[i]], j);
    }
    sort(all(dist));
    // cout<<dist.size()<<" "<<out[b[i]]-in[b[i]]<<endl;
    rep(j, min((int)dist.size(), (out[b[i]] - in[b[i]]) * 10)) {
      // rep(j, dist.size()) {
      graph.add_edge(dist[j].second, a.size() + i, inf, dist[j].first);
      esz++;
    }
  }
  //*/
  // cout << esz << endl;
  // auto start = chrono::system_clock::now();
  printf("%lld\\n", sum + graph.mincost_flow(s, t, outnum));
  // auto end = chrono::system_clock::now();
  // cout << fixed << setprecision(12) <<
  // chrono::duration_cast<chrono::seconds>(end-start).count() << endl;

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;

#define int long long
#define all(v) (v).begin(), (v).end()
#define resz(v, ...) (v).clear(), (v).resize(__VA_ARGS__)
#define reps(i, m, n) for (int i = (int)(m); i < (int)(n); i++)
#define rep(i, n) reps(i, 0, n)

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

using Pi = pair<int, int>;
using Tapris = tuple<int, int, int>;
using vint = vector<int>;

const int inf = 1LL << 55;
const int mod = 1e9 + 7;

struct edge {
  int to, cap, cost, rev;
  edge() {}
  edge(int to, int cap, int cost, int rev)
      : to(to), cap(cap), cost(cost), rev(rev) {}
};

using Graph = vector<vector<edge>>;

struct PrimalDual {
  Graph graph;
  vint mincost, h, prevv, preve;
  PrimalDual(int n) : graph(n), mincost(n), h(n), prevv(n), preve(n) {}
  void add_edge(int from, int to, int cap, int cost) {
    graph[from].emplace_back(to, cap, cost, graph[to].size());
    graph[to].emplace_back(from, 0, -cost, graph[from].size() - 1);
  }
  int mincost_flow(int s, int t, int f) {
    fill(all(h), 0);
    fill(all(prevv), -1);
    fill(all(preve), -1);
    int res = 0;
    while (f > 0) {
      fill(all(mincost), inf);
      mincost[s] = 0;
      priority_queue<Pi, vector<Pi>, greater<Pi>> que;
      que.emplace(0, s);
      while (!que.empty()) {
        int cost, v;
        tie(cost, v) = que.top();
        que.pop();
        if (mincost[v] < cost)
          continue;
        for (int i = 0; i < (int)graph[v].size(); ++i) {
          edge &e = graph[v][i];
          int dual = cost + e.cost + h[v] - h[e.to];
          if (e.cap > 0 && dual < mincost[e.to]) {
            mincost[e.to] = dual;
            prevv[e.to] = v;
            preve[e.to] = i;
            que.emplace(dual, e.to);
          }
        }
      }
      if (mincost[t] == inf)
        return res;
      for (int v = 0; v < (int)graph.size(); ++v)
        h[v] += mincost[v];
      int d = f;
      for (int v = t; v != s; v = prevv[v])
        chmin(d, graph[prevv[v]][preve[v]].cap);
      f -= d;
      res += d * h[t];
      for (int v = t; v != s; v = prevv[v]) {
        edge &e = graph[prevv[v]][preve[v]];
        e.cap -= d;
        graph[e.to][e.rev].cap += d;
      }
    }
    return res;
  }
};

signed main() {
  cin.tie(0);
  ios_base::sync_with_stdio(0);
  cout << fixed << setprecision(12);

  int n;
  cin >> n;
  vector<vint> dp(n, vint(n, inf));
  vint in(n, 0), out(n, 0);
  int sum = 0;
  rep(i, n) dp[i][i] = 0;
  in[0] = inf;
  rep(i, n) {
    int k;
    cin >> k;
    rep(j, k) {
      int t, c;
      cin >> t >> c;
      --t;
      chmin(dp[i][t], c);
      sum += c;
      ++in[t];
      ++out[i];
    }
  }
  rep(i, n) rep(j, n) rep(k, n) chmin(dp[j][k], dp[j][i] + dp[i][k]);
  vint a, b;
  rep(i, n) {
    if (in[i] > out[i])
      a.push_back(i);
    else if (in[i] < out[i])
      b.push_back(i);
  }
  int s = a.size() + b.size(), t = s + 1, sz = t + 1;
  PrimalDual graph(sz);
  rep(i, a.size()) graph.add_edge(s, i, in[a[i]] - out[a[i]], 0);
  rep(i, b.size()) graph.add_edge(a.size() + i, t, out[b[i]] - in[b[i]], 0);
  rep(i, a.size()) {
    rep(j, b.size()) {
      if (dp[a[i]][b[j]] < inf) {
        graph.add_edge(i, a.size() + j, inf, dp[a[i]][b[j]]);
      }
    }
  }
  cout << sum + graph.mincost_flow(s, t, inf) << endl;

  return 0;
}