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

#include <algorithm>
#include <chrono>
#include <cmath>
#include <deque>
#include <fstream>
#include <functional>
#include <iostream>
#include <iterator>
#include <limits>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#define forr(i, start, count) for (int i = (start); i < (start) + (count); ++i)
#define set_map_includes(set, elt) (set.find((elt)) != set.end())
#define readint(i)                                                             \\
  int i;                                                                       \\
  cin >> i
#define readll(i)                                                              \\
  ll i;                                                                        \\
  cin >> i
#define readdouble(i)                                                          \\
  double i;                                                                    \\
  cin >> i
#define readstring(s)                                                          \\
  string s;                                                                    \\
  cin >> s

typedef long long ll;

using namespace std;

ll modd = 1000 * 1000 * 1000 + 7;

template <class T> class MinCostMaxFlow {
  // between two vertices there can only be one directed edge!!! if you want the
  // trick with several edges between vertices,  break down the edge into two
  // pieces by introduction of new vertex costs can be negative!! (capacities
  // ofcourse not)
public:
  struct Edge {
    int from, to;
    T capacity, cost;
  };
  vector<vector<int>> adj;
  vector<vector<T>> cost, capacity;
  vector<Edge> edges;
  const ll INF = modd * modd;

  void add_edge(int fr, int tt, T cap, T cst) {
    edges.push_back({fr, tt, cap, cst});
  }

  void clear_edges() { edges.clear(); }

  void shortest_paths(int n, int v0, vector<T> &d, vector<T> &p) {
    d.assign(n, INF);
    d[v0] = 0;
    vector<bool> inq(n, false);
    queue<int> q;
    q.push(v0);
    p.assign(n, -1);

    while (!q.empty()) {
      int u = q.front();
      q.pop();
      inq[u] = false;
      for (int v : adj[u]) {
        if (capacity[u][v] > 0 && d[v] > d[u] + cost[u][v]) {
          d[v] = d[u] + cost[u][v];
          p[v] = u;
          if (!inq[v]) {
            inq[v] = true;
            q.push(v);
          }
        }
      }
    }
  }

  int flow_opt;

  pair<T, map<pair<int, int>, T>> min_cost_flow(int N, T K, int s, int t) {
    // looks for min cost flow among all flows of total flow less or equal K;
    // return min cost flow among flows of total flow K, or if such flow doesnt
    // exist, min cost flow for max flow if you need the value of output flow,
    // it's stored in variable flow at the time of return if you need min cost
    // flow among all flows <= K, take minimum of flow where it says here!!!!!
    // and store it and output it, bc cost can decrease in next iteration when
    // flow increases if you need min cost max flow, take very large K
    adj.assign(N, vector<int>());
    cost.assign(N, vector<T>(N, 0));
    capacity.assign(N, vector<T>(N, 0));
    map<pair<int, int>, T> ret;
    for (Edge e : edges) {
      adj[e.from].push_back(e.to);
      adj[e.to].push_back(e.from);
      cost[e.from][e.to] = e.cost;
      cost[e.to][e.from] = -e.cost;
      capacity[e.from][e.to] = e.capacity;
    }

    T flow = 0;
    T cost = 0;
    vector<T> d, p;
    T min_cost = INF;
    flow_opt = 0;
    while (flow < K) {
      shortest_paths(N, s, d, p);
      if (d[t] == INF)
        break;

      T f = K - flow;
      int cur = t;
      while (cur != s) {
        f = min(f, capacity[p[cur]][cur]);
        cur = p[cur];
      }

      // apply flow
      flow += f;
      cost += f * d[t];
      cur = t;
      if (cost < min_cost) {
        min_cost = cost;
        flow_opt = flow;
        //            cout << min_cost << " " << flow_opt << endl;
      }
      while (cur != s) {
        ret[make_pair(p[cur], cur)] += f;
        capacity[p[cur]][cur] -= f;
        capacity[cur][p[cur]] += f;
        cur = p[cur];
      }
    }
    for (Edge e : edges) {
      if (set_map_includes(ret, make_pair(e.to, e.from))) {
        ret[make_pair(e.from, e.to)] -= ret[make_pair(e.to, e.from)];
        ret.erase(make_pair(e.to, e.from));
      }
    }

    return make_pair(cost, ret);
  }
};

int main() {

  ios_base::sync_with_stdio(false);

  cout.precision(17);
  mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
  uniform_int_distribution<int> rand_gen(
      0, modd); // rand_gen(rng) gets the rand no

  //    readint(test_cases);
  int test_cases = 1;
  forr(t, 1, test_cases) {

    readint(n);
    readint(k);
    vector<vector<ll>> a;
    forr(i, 0, n) {
      a.push_back(vector<ll>());
      forr(j, 0, n) {
        readint(aa);
        a.back().push_back(aa);
      }
    }
    MinCostMaxFlow<ll> mflow;

    forr(i, 0, n) {
      mflow.add_edge(2 * n, i, k, 0);
      mflow.add_edge(i + n, 2 * n + 1, k, 0);
    }
    forr(i, 0, n) {
      forr(j, 0, n) { mflow.add_edge(i, j + n, 1, -a[i][j]); }
    }
    auto q = mflow.min_cost_flow(2 * n + 2, n * k, 2 * n, 2 * n + 1);

    int uuu = mflow.flow_opt;

    mflow.clear_edges();
    forr(i, 0, n) {
      mflow.add_edge(2 * n, i, k, 0);
      mflow.add_edge(i + n, 2 * n + 1, k, 0);
    }
    forr(i, 0, n) {
      forr(j, 0, n) { mflow.add_edge(i, j + n, 1, -a[i][j]); }
    }
    q = mflow.min_cost_flow(2 * n + 2, uuu, 2 * n, 2 * n + 1);
    cout << -q.first << endl;
    forr(i, 0, n) {
      forr(j, 0, n) {
        cout << (q.second[make_pair(i, j + n)] == 1 ? 'X' : '.');
      }
      cout << endl;
    }
  }
  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <chrono>
#include <cmath>
#include <deque>
#include <fstream>
#include <functional>
#include <iostream>
#include <iterator>
#include <limits>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <sstream>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#define forr(i, start, count) for (int i = (start); i < (start) + (count); ++i)
#define set_map_includes(set, elt) (set.find((elt)) != set.end())
#define readint(i)                                                             \\
  int i;                                                                       \\
  cin >> i
#define readll(i)                                                              \\
  ll i;                                                                        \\
  cin >> i
#define readdouble(i)                                                          \\
  double i;                                                                    \\
  cin >> i
#define readstring(s)                                                          \\
  string s;                                                                    \\
  cin >> s

typedef long long ll;

using namespace std;

ll modd = 1000 * 1000 * 1000 + 7;

template <class T> class MinCostMaxFlow {
  // between two vertices there can only be one directed edge!!! if you want the
  // trick with several edges between vertices,  break down the edge into two
  // pieces by introduction of new vertex costs can be negative!! (capacities
  // ofcourse not)
public:
  struct Edge {
    int from, to;
    T capacity, cost;
  };
  vector<vector<int>> adj;
  vector<vector<T>> cost, capacity;
  vector<Edge> edges;
  const ll INF = modd * modd;

  void add_edge(int fr, int tt, T cap, T cst) {
    edges.push_back({fr, tt, cap, cst});
  }

  void clear_edges() { edges.clear(); }

  void shortest_paths(int n, int v0, vector<T> &d, vector<T> &p) {
    d.assign(n, INF);
    d[v0] = 0;
    vector<bool> inq(n, false);
    queue<int> q;
    q.push(v0);
    p.assign(n, -1);

    while (!q.empty()) {
      int u = q.front();
      q.pop();
      inq[u] = false;
      for (int v : adj[u]) {
        if (capacity[u][v] > 0 && d[v] > d[u] + cost[u][v]) {
          d[v] = d[u] + cost[u][v];
          p[v] = u;
          if (!inq[v]) {
            inq[v] = true;
            q.push(v);
          }
        }
      }
    }
  }

  pair<T, map<pair<int, int>, T>> min_cost_flow(int N, T K, int s, int t) {
    adj.assign(N, vector<int>());
    cost.assign(N, vector<T>(N, 0));
    capacity.assign(N, vector<T>(N, 0));
    map<pair<int, int>, T> ret;
    for (Edge e : edges) {
      adj[e.from].push_back(e.to);
      adj[e.to].push_back(e.from);
      cost[e.from][e.to] = e.cost;
      cost[e.to][e.from] = -e.cost;
      capacity[e.from][e.to] = e.capacity;
    }

    T flow = 0;
    T cost = 0;
    vector<T> d, p;
    while (flow < K) {
      shortest_paths(N, s, d, p);
      if (d[t] == INF)
        break;

      T f = K - flow;
      int cur = t;
      while (cur != s) {
        f = min(f, capacity[p[cur]][cur]);
        cur = p[cur];
      }

      // apply flow
      flow += f;
      cost += f * d[t];
      cur = t;
      while (cur != s) {
        ret[make_pair(p[cur], cur)] += f;
        capacity[p[cur]][cur] -= f;
        capacity[cur][p[cur]] += f;
        cur = p[cur];
      }
    }
    for (Edge e : edges) {
      if (set_map_includes(ret, make_pair(e.to, e.from))) {
        ret[make_pair(e.from, e.to)] -= ret[make_pair(e.to, e.from)];
        ret.erase(make_pair(e.to, e.from));
      }
    }

    if (flow < K)
      return make_pair(-1, ret);
    else
      return make_pair(cost, ret);
  }
};

int main() {

  ios_base::sync_with_stdio(false);

  cout.precision(17);
  mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
  uniform_int_distribution<int> rand_gen(
      0, modd); // rand_gen(rng) gets the rand no

  //    readint(test_cases);
  int test_cases = 1;
  forr(t, 1, test_cases) {
    readint(n);
    readint(k);
    vector<vector<ll>> a;
    forr(i, 0, n) {
      a.push_back(vector<ll>());
      forr(j, 0, n) {
        readint(aa);
        a.back().push_back(aa);
      }
    }
    MinCostMaxFlow<ll> mflow;

    ll maxfl = 0;
    int p_max;
    forr(p, 1, k * n) {
      mflow.clear_edges();
      forr(i, 0, n) {
        mflow.add_edge(2 * n, i, k, 0);
        mflow.add_edge(i + n, 2 * n + 1, k, 0);
      }
      forr(i, 0, n) {
        forr(j, 0, n) { mflow.add_edge(i, j + n, 1, -a[i][j]); }
      }
      ll vv = mflow.min_cost_flow(2 * n + 2, p, 2 * n, 2 * n + 1).first;
      if (vv < maxfl) {
        maxfl = vv;
        p_max = p;
      }
    }

    mflow.clear_edges();
    forr(i, 0, n) {
      mflow.add_edge(2 * n, i, k, 0);
      mflow.add_edge(i + n, 2 * n + 1, k, 0);
    }
    forr(i, 0, n) {
      forr(j, 0, n) { mflow.add_edge(i, j + n, 1, -a[i][j]); }
    }
    auto q = mflow.min_cost_flow(2 * n + 2, p_max, 2 * n, 2 * n + 1);
    cout << -q.first << endl;
    forr(i, 0, n) {
      forr(j, 0, n) {
        cout << (q.second[make_pair(i, j + n)] == 1 ? 'X' : '.');
      }
      cout << endl;
    }
  }
  return 0;
}