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

#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (ll i = 0; i < n; i++)
#define rep2(i, x, n) for (ll i = x; i <= n; i++)
#define rep3(i, x, n) for (int i = x; i >= n; i--)
#define elif else if
#define sp(x) fixed << setprecision(x)
#define pb(x) push_back(x)
#define all(x) x.begin(), x.end()
using ll = long long;
using ld = long double;
using pii = pair<int, int>;
using pil = pair<int, ll>;
using pli = pair<ll, int>;
using pll = pair<ll, ll>;
const ll MOD = 1e9 + 7;
// const ll MOD = 998244353;
const int inf = 1e9;
const ll INF = 1e18;
const ld EPS = 1e-10;

int main() {
  int N;
  cin >> N;
  pli p[N];
  rep(i, N) {
    cin >> p[i].first;
    p[i].second = i;
  }
  sort(p, p + N, greater<pli>());
  // dp[i][j]:=[0,i)と[j,N)まで決まっているときの最大値
  ll dp[N + 1][N + 1];
  fill(dp[0], dp[N + 1], 0);
  rep(i, N) {
    ll A = p[i].first, id = p[i].second;
    // すでにi個並べている
    //[0,j)と[N+j-i,N)が埋まっているとき
    rep(j, i + 1) {
      dp[j + 1][N + j - i] =
          max(dp[j + 1][N + j - i], dp[j][N + j - i] + A * abs(j - id));
      dp[j][N + j - i - 1] = max(
          dp[j][N + j - i - 1], dp[j][N + j - i] + A * abs(N + j - i - 1 - id));
    }
  }
  ll ans = 0;
  rep(i, N + 1) { ans = max(ans, dp[i][i]); }
  cout << ans << endl;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define rep(i, n) for (int i = 0; i < n; i++)
#define rep2(i, x, n) for (int i = x; i <= n; i++)
#define rep3(i, x, n) for (int i = x; i >= n; i--)
#define elif else if
#define sp(x) fixed << setprecision(x)
#define pb(x) push_back(x)
#define all(x) x.begin(), x.end()
using ll = long long;
using ld = long double;
using pii = pair<int, int>;
using pil = pair<int, ll>;
using pli = pair<ll, int>;
using pll = pair<ll, ll>;
const ll MOD = 1e9 + 7;
// const ll MOD = 998244353;
const int inf = 1e9;
const ll INF = 1e18;
const ld EPS = 1e-10;

struct edge {
  int to;
  ll cap, cost;
  int rev;
};

int V, E;
const int MAX_V = 5e3;
vector<edge> es[MAX_V];
ll d[MAX_V], h[MAX_V];
int pre_e[MAX_V], pre_v[MAX_V];

void add_edge(int from, int to, ll cap, ll cost) {
  edge e1 = {to, cap, cost, (int)es[to].size()};
  edge e2 = {from, 0, -cost, (int)es[from].size()};
  es[from].emplace_back(e1), es[to].emplace_back(e2);
}

void dijkstra(int s) {
  fill(d, d + V, INF);
  priority_queue<pli, vector<pli>, greater<pli>> que;
  d[s] = 0;
  que.push(pli(0, s));
  while (!que.empty()) {
    pli p = que.top();
    que.pop();
    int i = p.second;
    if (d[i] < p.first)
      continue;
    rep(j, es[i].size()) {
      edge &e = es[i][j];
      if (e.cap > 0 && d[e.to] > d[i] + e.cost + h[i] - h[e.to]) {
        d[e.to] = d[i] + e.cost + h[i] - h[e.to];
        pre_v[e.to] = i;
        pre_e[e.to] = j;
        que.push(pli(d[e.to], e.to));
      }
    }
  }
}

// sからtへの流量flowの最小費用流(流せない場合は-1)を最大でO(F*E*log(V))で計算
ll min_cost_flow(int s, int t, ll flow) {
  ll res = 0;
  fill(h, h + V, 0);
  while (flow > 0) {
    dijkstra(s);
    if (d[t] == INF)
      return -1;
    rep(i, V) h[i] += d[i];
    ll f = flow;
    for (int now = t; now != s; now = pre_v[now]) {
      f = min(f, es[pre_v[now]][pre_e[now]].cap);
    }
    flow -= f;
    res += f * h[t];
    for (int now = t; now != s; now = pre_v[now]) {
      edge &e = es[pre_v[now]][pre_e[now]];
      e.cap -= f;
      es[now][e.rev].cap += f;
    }
  }
  return res;
}

int main() {
  int N;
  cin >> N;
  V = 2 * N + 2;
  ll A[N], sum = 0;
  rep(i, N) cin >> A[i];
  rep(i, N) sum += A[i];
  rep(i, N) {
    add_edge(2 * N, i, 1, 0);
    add_edge(N + i, 2 * N + 1, 1, 0);
    rep(j, N) { add_edge(i, N + j, 1, A[i] * (N - abs(j - i))); }
  }
  cout << N * sum - min_cost_flow(2 * N, 2 * N + 1, N) << endl;
}