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

#pragma GCC target("avx2")
#pragma GCC optimize("unroll-loops")
#pragma GCC optimize("O3")
#include "bits/stdc++.h"
#include <random>
#include <unordered_map>
#include <unordered_set>
using namespace std;
typedef long long ll;
const ll MOD = 1'000'000'007LL; /*998'244'353LL;*/
#define pb push_back
#define mp make_pair
#define all(x) (x).begin(), (x).end()
#define rep(i, n) for (int(i) = 0; (i) < (n); (i)++)
const int dx[4] = {1, 0, -1, 0};
const int dy[4] = {0, 1, 0, -1};

struct MinCostFlow {
private:
  typedef pair<int, int> P;
  struct edge {
    int to, cap;
    long long cost;
    int rev, num;
    bool isrev;
  };

  int V;
  vector<vector<edge>> G;
  vector<long long> h, dist;
  vector<int> prevv, preve;
  const long long INF = 1LL << 62;

public:
  vector<long long> edges;

  MinCostFlow(int _V = 0) {
    V = _V;
    G.resize(V);
    h.resize(V);
    dist.resize(V);
    prevv.resize(V);
    preve.resize(V);
  }

  void add(int from, int to, int cap, long long cost) {
    G[from].push_back(
        {to, cap, cost, (int)G[to].size(), (int)edges.size(), false});
    G[to].push_back(
        {from, 0, -cost, (int)G[from].size() - 1, (int)edges.size(), true});
    edges.push_back(0);
  }

  long long calculate(int s, int t, int f = INT_MAX) {
    long long ret = 0;
    h.assign(V, 0);
    while (f > 0) {
      priority_queue<P, vector<P>, greater<P>> que;
      dist.assign(V, INF);
      dist[s] = 0;
      que.push({0, 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 < G[v].size(); i++) {
          edge &e = G[v][i];
          if (e.cap > 0 && 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.push({dist[e.to], e.to});
          }
        }
      }
      if (dist[t] == INF)
        break;
      for (int v = 0; v < V; v++)
        h[v] += dist[v];

      int d = f;
      for (int v = t; v != s; v = prevv[v]) {
        d = min(d, G[prevv[v]][preve[v]].cap);
      }
      f -= d;
      ret += 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;
        if (!e.isrev)
          edges[e.num] += d;
        else
          edges[e.num] -= d;
      }
    }
    return ret;
  }

  vector<pair<int, long long>> calculate_slope(int s, int t) {
    int f = 0;
    long long cost = 0;
    vector<pair<int, long long>> ret;
    h.assign(V, 0);
    while (true) {
      priority_queue<P, vector<P>, greater<P>> que;
      dist.assign(V, INF);
      dist[s] = 0;
      que.push({0, 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 < G[v].size(); i++) {
          edge &e = G[v][i];
          if (e.cap > 0 && 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.push({dist[e.to], e.to});
          }
        }
      }
      if (dist[t] == INF)
        break;
      for (int v = 0; v < V; v++)
        h[v] += dist[v];

      int d = INT_MAX;
      for (int v = t; v != s; v = prevv[v]) {
        d = min(d, G[prevv[v]][preve[v]].cap);
      }
      f += d;
      cost += d * h[t];
      ret.push_back({f, cost});
      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;
        if (!e.isrev)
          edges[e.num] += d;
        else
          edges[e.num] -= d;
      }
    }
    return ret;
  }
};

int N, K;
int A[50][50];

signed main() {
  cin >> N >> K;
  rep(i, N) {
    rep(j, N) { cin >> A[i][j]; }
  }

  MinCostFlow f(1 + N + N * N + N + 1);
  rep(i, N) { rep(j, N) f.add(1 + i, 1 + N + i * N + j, 1, -A[i][j]); }
  rep(i, N) { rep(j, N) f.add(1 + N + i * N + j, 1 + N + N * N + j, 1, 0); }
  rep(i, N) f.add(0, 1 + i, K, 0);
  rep(j, N) f.add(1 + N + N * N + j, 1 + N + N * N + N, K, 0);

  MinCostFlow tmp = f;
  auto v = tmp.calculate_slope(0, 1 + N + N * N + N);
  long long bef = LLONG_MAX;
  int idx = v.back().first;
  rep(i, v.size()) {
    if (bef < v[i].second) {
      idx = v[i - 1].first;
      break;
    }
    bef = v[i].second;
  }

  cout << -f.calculate(0, 1 + N + N * N + N, idx) << endl;
  rep(i, N) {
    rep(j, N) {
      if (f.edges[i * N + j] > 0)
        cout << "X";
      else
        cout << ".";
    }
    cout << endl;
  }
}

Error: TLE

Fixed submission:

#pragma GCC target("avx2")
#pragma GCC optimize("unroll-loops")
#pragma GCC optimize("O3")
#include "bits/stdc++.h"
#include <random>
#include <unordered_map>
#include <unordered_set>
using namespace std;
typedef long long ll;
const ll MOD = 1'000'000'007LL; /*998'244'353LL;*/
#define pb push_back
#define mp make_pair
#define all(x) (x).begin(), (x).end()
#define rep(i, n) for (int(i) = 0; (i) < (n); (i)++)
const int dx[4] = {1, 0, -1, 0};
const int dy[4] = {0, 1, 0, -1};

struct MinCostFlow {
private:
  typedef pair<int, int> P;
  struct edge {
    int to, cap;
    long long cost;
    int rev, num;
    bool isrev;
  };

  int V;
  vector<vector<edge>> G;
  vector<long long> h, dist;
  vector<int> prevv, preve;
  const long long INF = 1LL << 62;

public:
  vector<long long> edges;

  MinCostFlow(int _V = 0) {
    V = _V;
    G.resize(V);
    h.resize(V);
    dist.resize(V);
    prevv.resize(V);
    preve.resize(V);
  }

  void add(int from, int to, int cap, long long cost) {
    G[from].push_back(
        {to, cap, cost, (int)G[to].size(), (int)edges.size(), false});
    G[to].push_back(
        {from, 0, -cost, (int)G[from].size() - 1, (int)edges.size(), true});
    edges.push_back(0);
  }

  long long calculate(int s, int t, int f = INT_MAX) {
    long long ret = 0;
    h.assign(V, 0);
    while (f > 0) {
      priority_queue<P, vector<P>, greater<P>> que;
      dist.assign(V, INF);
      dist[s] = 0;
      que.push({0, 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 < G[v].size(); i++) {
          edge &e = G[v][i];
          if (e.cap > 0 && 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.push({dist[e.to], e.to});
          }
        }
      }
      if (dist[t] == INF)
        break;
      for (int v = 0; v < V; v++)
        h[v] += dist[v];

      int d = f;
      for (int v = t; v != s; v = prevv[v]) {
        d = min(d, G[prevv[v]][preve[v]].cap);
      }
      f -= d;
      ret += 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;
        if (!e.isrev)
          edges[e.num] += d;
        else
          edges[e.num] -= d;
      }
    }
    return ret;
  }
};

int N, K;
int A[50][50];

signed main() {
  cin >> N >> K;
  rep(i, N) {
    rep(j, N) { cin >> A[i][j]; }
  }

  MinCostFlow f(1 + N + N * N + N + 1);
  rep(i, N) { rep(j, N) f.add(1 + i, 1 + N + i * N + j, 1, -A[i][j]); }
  rep(i, N) { rep(j, N) f.add(1 + N + i * N + j, 1 + N + N * N + j, 1, 0); }
  rep(i, N) f.add(0, 1 + i, K, 0);
  rep(j, N) f.add(1 + N + N * N + j, 1 + N + N * N + N, K, 0);

  int inc = 0, dec = K * N;
  while (dec - inc > 1) {
    int m = (inc + dec) / 2;
    MinCostFlow tmp1 = f, tmp2 = f;
    if (-tmp1.calculate(0, 1 + N + N * N + N, m) <
        -tmp2.calculate(0, 1 + N + N * N + N, m + 1))
      inc = m;
    else
      dec = m;
  }
  cout << -f.calculate(0, 1 + N + N * N + N, dec) << endl;
  rep(i, N) {
    rep(j, N) {
      if (f.edges[i * N + j] > 0)
        cout << "X";
      else
        cout << ".";
    }
    cout << endl;
  }
}