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

#include <algorithm>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <utility>
#include <vector>

using namespace std;

template <class T> class Min_Cost_Flow {
public:
  struct edge {
    int to;
    int cap;
    T cost;
    int rev;
    bool is_rev;
  };

  const T INF;
  vector<vector<edge>> G;

  Min_Cost_Flow(int n, T inf) : G(n), INF(inf) {}

  void add_edge(int from, int to, int cap, T cost) {
    G[from].push_back((edge){to, cap, cost, (int)G[to].size(), false});
    G[to].push_back((edge){from, 0, -cost, (int)G[from].size() - 1, true});
  }

  // min cost : s->t (flow:f)
  T min_cost_flow(int s, int t, int f) {
    T cost = 0;
    vector<int> prev_v(G.size(), -1);
    vector<int> prev_e(G.size(), -1);
    vector<T> potantial(G.size(), 0);
    vector<T> dist(G.size(), INF);

    dist[s] = 0;

    /*bellman_ford*/ {
      bool update = true;
      for (int cnt = 0; update; cnt++) {
        update = false;
        for (int i = 0; i < G.size(); i++) {
          if (dist[i] >= INF)
            continue;
          for (int j = 0; j < G[i].size(); j++) {
            if (G[i][j].cap > 0 && dist[G[i][j].to] > dist[i] + G[i][j].cost) {
              dist[G[i][j].to] = dist[i] + G[i][j].cost;

              update = true;
            }
          }
        }

        if (update && cnt >= G.size()) {
          cerr << " there is a negative cycle " << endl;
          return -1;
        }
        cnt++;
      }

      for (int i = 0; i < G.size(); i++) {
        potantial[i] = dist[i];
      }
    }

    while (f > 0) {
      fill(dist.begin(), dist.end(), INF);
      dist[s] = 0;
      prev_v[s] = s;

      /*dijkstra*/ {
        priority_queue<pair<T, int>, vector<pair<T, int>>,
                       greater<pair<T, int>>>
            pq;
        pq.push({0, s});
        while (pq.size()) {
          pair<T, int> p = pq.top();
          pq.pop();
          if (dist[p.second] < p.first)
            continue;
          dist[p.second] = p.first;
          for (int i = 0; i < G[p.second].size(); i++) {
            edge &e = G[p.second][i];
            T new_dist =
                dist[p.second] + e.cost + potantial[p.second] - potantial[e.to];
            if (e.cap > 0 && dist[e.to] > new_dist) {
              dist[e.to] = new_dist;
              prev_v[e.to] = p.second;
              prev_e[e.to] = i;
              pq.push({dist[e.to], e.to});
            }
          }
        }
      }

      if (potantial[t] + dist[t] >= INF) {
        cerr << "couldn't insert flow f : " << f << endl;
        return -1;
      }

      for (int i = 0; i < G.size(); i++) {
        potantial[i] += dist[i];
      }

      int d = f;
      int pos = t;
      while (pos != s) {
        int v = prev_v[pos];
        int e = prev_e[pos];
        d = min(d, G[v][e].cap);
        pos = prev_v[pos];
      }

      pos = t;
      while (pos != s) {
        int v = prev_v[pos];
        int e = prev_e[pos];
        G[v][e].cap -= d;
        G[G[v][e].to][G[v][e].rev].cap += d;
        pos = v;
      }
      cost += d * potantial[t]; // potential is new_dist
      f -= d;

      // f==0 then end
    }
    return cost;
  }
};

long long solve(int n, vector<long long> a) {
  Min_Cost_Flow<long long> f(2 * n + 2, (1ll << 58));
  int source = n * 2;
  int sink = source + 1;
  for (int i = 0; i < n; i++) {
    f.add_edge(source, i, 1, 0);
    f.add_edge(n + i, sink, 1, 0);
    for (int j = 0; j < n; j++) {
      long long d = abs(i - j);
      f.add_edge(i, n + j, 1, -d * a[i]);
    }
  }

  long long ans = f.min_cost_flow(source, sink, n);

  vector<int> res(n);
  for (int i = 0; i < n; i++) {
    for (auto e : f.G[i]) {
      if (e.is_rev)
        continue;
      if (e.cap != 0)
        continue;
      res[e.to - n] = i;
    }
  }
  // for(int i=0; i<n; i++){
  //   cerr << res[i] << " ";
  // }
  // cerr << endl;
  // for(int i=0; i<n; i++){
  //   cerr << a[res[i]] << " ";
  // }
  // cerr << endl;

  return -ans;
}

long long solve2(int n, vector<long long> a) {
  vector<vector<long long>> dp(n + 1, vector<long long>(n + 1, -(1ll << 58)));
  vector<int> v(n);
  iota(begin(v), end(v), 0);
  sort(begin(v), end(v), [&](int i, int j) {
    if (a[i] != a[j])
      return a[i] > a[j];
    return max(i, n - 1 - i) > max(j, n - 1 - j);
  });

  dp[0][0] = 0;
  for (int i = 0; i < n; i++) {
    int j = v[i];
    for (int l = 0; l <= i; l++) {
      int r = i - l;
      int li = l;
      int ri = n - 1 - r;

      {
        long long d = abs(j - li);
        long long tmp = d * a[j];
        dp[i + 1][l + 1] = max(dp[i + 1][l + 1], dp[i][l] + tmp);
      }
      {
        long long d = abs(j - ri);
        long long tmp = d * a[j];
        dp[i + 1][l] = max(dp[i + 1][l], dp[i][l] + tmp);
      }
    }
  }

  long long ans = *max_element(begin(dp[n]), end(dp[n]));
  return ans;
}

#include <cassert>
#include <random>
int main() {
  int n;
  cin >> n;
  vector<long long> a(n);
  for (int i = 0; i < n; i++) {
    cin >> a[i];
  }

  // mt19937 mt(19937);
  // for(int itr=0; itr<100000; itr++){
  //   int n = mt()%20 + 1;
  //   vector<long long> a(n);
  //   for(int i=0; i<n; i++){
  //     a[i] = mt()%100 + 1;
  //   }
  //   long long a0 = solve(n, a);
  //   // cerr << a0 << endl;
  //   long long a1 = solve2(n, a);
  //   if(a0 != a1){
  //     cerr << n << endl;
  //     for(int i=0; i<n; i++){
  //       cerr << a[i] << " ";
  //     }
  //     cerr << endl;
  //     cerr << a0 << endl;
  //     cerr << a1 << endl;
  //     assert(a0 == a1);
  //   }
  // }

  // long long ans = solve(n, a);
  long long ans2 = solve2(n, a);
  // cout << ans << endl;
  cout << ans2 << endl;
  return 0;
}

Error: TLE

Fixed submission:

#include <algorithm>
#include <iostream>
#include <map>
#include <queue>
#include <vector>

using namespace std;

template <class T> class Min_Cost_Flow {
public:
  struct edge {
    int to;
    int cap;
    T cost;
    int rev;
  };

  const T INF;
  vector<vector<edge>> G;

  Min_Cost_Flow(int n, T inf) : G(n), INF(inf) {}

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

  // min cost : s->t (flow:f)
  T min_cost_flow(int s, int t, int f) {
    const int N = G.size();
    T cost = 0;
    vector<int> prev_v(N, -1);
    vector<int> prev_e(N, -1);
    vector<T> potantial(N, 0);

    while (f > 0) {
      // min distance(cost based) search with SPFA
      vector<T> dist(N, INF);

      vector<int> cnt(dist.size(), 0);
      dist[s] = 0;
      prev_v[s] = s;

      queue<int> Q;
      auto my_push = [&](int node) {
        Q.push(node);
        cnt[node]++;
      };
      auto my_pop = [&]() -> int {
        int ret = Q.front();
        cnt[ret]--;
        Q.pop();
        return ret;
      };
      my_push(s);
      while (!Q.empty()) {
        int pos = my_pop();
        for (int i = 0; i < G[pos].size(); i++) {
          edge &E = G[pos][i];
          T new_dist = dist[pos] + E.cost + potantial[E.to] - potantial[pos];
          if (dist[E.to] > new_dist && E.cap > 0) {
            dist[E.to] = new_dist;
            prev_v[E.to] = pos;
            prev_e[E.to] = i;

            if (cnt[E.to] == 0) {
              my_push(E.to);
            }
          }
        }
      }
      for (int i = 0; i < N; i++) {
        dist[i] = potantial[i] + dist[i];
      }

      // cannot achieved to "t" return -1
      if (dist[t] >= INF)
        return -1;

      // add cost of s->t with flow=d
      int pos = t;
      int d = f;
      while (pos != s) {
        int i = prev_v[pos];
        int j = prev_e[pos];
        pos = i;
        d = min(d, G[i][j].cap);
      }

      pos = t;
      // cout << t ;
      while (pos != s) {
        int i = prev_v[pos];
        int j = prev_e[pos];
        G[i][j].cap -= d;
        G[G[i][j].to][G[i][j].rev].cap += d;
        // cost += G[i][j].cost * d;
        pos = i;
        // cout << " <- " << pos;
      }
      // cout << endl;
      cost += d * dist[t];
      f -= d;

      // f==0 then end
    }
    return cost;
  }
};

int main() {
  int n;
  cin >> n;
  vector<long long> a(n);
  for (int i = 0; i < n; i++) {
    cin >> a[i];
  }

  Min_Cost_Flow<long long> f(2 * n + 2, (1ll << 58));
  int source = n * 2;
  int sink = source + 1;
  for (int i = 0; i < n; i++) {
    f.add_edge(source, i, 1, 0);
    f.add_edge(n + i, sink, 1, 0);
    for (int j = 0; j < n; j++) {
      long long d = abs(i - j);
      f.add_edge(i, n + j, 1, -d * a[i]);
    }
  }

  long long ans = f.min_cost_flow(source, sink, n);

  cout << -ans << endl;
  return 0;
}