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

#include <bits/stdc++.h>
using namespace std;

#define rep(i, m, n) for (int(i) = (int)(m); i < (int)(n); i++)
#define rep2(i, m, n) for (int(i) = (int)(n)-1; i >= (int)(m); i--)
#define REP(i, n) rep(i, 0, n)
#define REP2(i, n) rep2(i, 0, n)
#define FOR(i, c)                                                              \\
  for (decltype((c).begin()) i = (c).begin(); i != (c).end(); ++i)
#define all(hoge) (hoge).begin(), (hoge).end()
#define en '\\n'
using ll = long long;
using ull = unsigned long long;
template <class T> using vec = vector<T>;
template <class T> using vvec = vector<vec<T>>;
typedef pair<ll, ll> P;
constexpr long long INF = 1LL << 60;
constexpr int INF_INT = 1 << 25;
constexpr long long MOD = (ll)1e9 + 7;
// constexpr long long MOD = 998244353LL;
typedef vector<ll> Array;
typedef vector<Array> Matrix;

template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

struct Edge {
  ll to, cap, rev;
  Edge(ll _to, ll _cap, ll _rev) {
    to = _to;
    cap = _cap;
    rev = _rev;
  }
};
using Edges = vector<Edge>;
using Graph = vector<Edges>;

void add_edge(Graph &G, ll from, ll to, ll cap, bool revFlag, ll revCap) {
  G[from].push_back(Edge(to, cap, (ll)G[to].size()));
  if (revFlag)
    G[to].push_back(Edge(from, revCap, (ll)G[from].size() - 1));
}

void solve() {
  ll n;
  cin >> n;
  Array a(n);
  REP(i, n) cin >> a[i];

  vec<P> b(n);
  REP(i, n) b[i] = {a[i], i};
  sort(all(b), greater<P>());

  Matrix dp(n + 1, Array(n + 1, 0)); // i番目まで見て、j個左に置く
  REP(i, n) {
    REP(j, i + 1) {
      ll p = b[i].first;
      ll x = b[i].second;
      if (x >= j)
        chmax(dp[i + 1][j + 1], dp[i][j] + p * abs(x - j));
      if (x <= n - 1 - i + j)
        chmax(dp[i + 1][j], dp[i][j] + p * abs(x - (n - 1 - i + j)));
    }
  }

  ll ans = 0;
  REP(i, n + 1) { chmax(ans, dp[n][i]); }
  cout << ans << en;
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);

  solve();
  // ll t;cin>>t;REP(i,t) solve();

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

#define rep(i, m, n) for (int(i) = (int)(m); i < (int)(n); i++)
#define rep2(i, m, n) for (int(i) = (int)(n)-1; i >= (int)(m); i--)
#define REP(i, n) rep(i, 0, n)
#define REP2(i, n) rep2(i, 0, n)
#define FOR(i, c)                                                              \\
  for (decltype((c).begin()) i = (c).begin(); i != (c).end(); ++i)
#define all(hoge) (hoge).begin(), (hoge).end()
#define en '\\n'
using ll = long long;
using ull = unsigned long long;
template <class T> using vec = vector<T>;
template <class T> using vvec = vector<vec<T>>;
typedef pair<ll, ll> P;
constexpr long long INF = 1LL << 60;
constexpr int INF_INT = 1 << 25;
// constexpr long long MOD = (ll) 1e9 + 7;
constexpr long long MOD = 998244353LL;
typedef vector<ll> Array;
typedef vector<Array> Matrix;

template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

typedef ll FLOW;        // フローを表す型、今回は int 型
typedef ll COST;        // コストを表す型、今回は int 型
const int MAX_V = 5050; // グラフの最大ノード数

// グラフの辺の構造体
struct Edge {
  ll rev, from, to;
  FLOW cap, icap;
  COST cost;
  Edge(ll r, ll f, ll t, FLOW ca, COST co)
      : rev(r), from(f), to(t), cap(ca), icap(ca), cost(co) {}
};

// グラフ構造体
struct Graph {
  int V;
  vector<Edge> list[MAX_V];

  Graph(int n = 0) : V(n) {
    for (int i = 0; i < MAX_V; ++i)
      list[i].clear();
  }
  void init(int n = 0) {
    V = n;
    for (int i = 0; i < MAX_V; ++i)
      list[i].clear();
  }
  void resize(int n = 0) { V = n; }
  void reset() {
    for (int i = 0; i < V; ++i)
      for (int j = 0; j < list[i].size(); ++j)
        list[i][j].cap = list[i][j].icap;
  }
  inline vector<Edge> &operator[](int i) { return list[i]; }

  Edge &redge(Edge &e) {
    if (e.from != e.to)
      return list[e.to][e.rev];
    else
      return list[e.to][e.rev + 1];
  }

  void addedge(ll from, ll to, FLOW cap, COST cost) {
    list[from].push_back(Edge((int)list[to].size(), from, to, cap, cost));
    list[to].push_back(Edge((int)list[from].size() - 1, to, from, 0, -cost));
  }
};

// 最小費用流を求める関数
COST MinCostFlow(Graph &G, int s, int t, FLOW inif) {
  COST dist[MAX_V];
  int prevv[MAX_V];
  int preve[MAX_V];

  COST res = 0;
  FLOW f = inif;
  while (f > 0) {
    fill(dist, dist + G.V, INF);
    dist[s] = 0;
    while (true) {
      bool update = false;
      for (int v = 0; v < G.V; ++v) {
        if (dist[v] == INF)
          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) {
            dist[e.to] = dist[v] + e.cost;
            prevv[e.to] = v;
            preve[e.to] = i;
            update = true;
          }
        }
      }
      if (!update)
        break;
    }

    if (dist[t] == INF)
      return 0;

    FLOW d = f;
    for (int v = t; v != s; v = prevv[v]) {
      d = min(d, G[prevv[v]][preve[v]].cap);
    }
    f -= d;
    res += dist[t] * d;
    for (int v = t; v != s; v = prevv[v]) {
      Edge &e = G[prevv[v]][preve[v]];
      Edge &re = G.redge(e);
      e.cap -= d;
      re.cap += d;
    }
  }
  return res;
}

void solve() {

  ll n;
  cin >> n;
  Array a(n);
  REP(i, n) cin >> a[i];

  // グラフの定義 (ノード数を引数に)
  Graph G(2 * n + 2); // +2 は S, T の分

  // グラフに枝を張っていく
  for (ll i = 0; i < n; ++i) {
    for (ll j = 0; j < n; ++j) {
      ll gain = a[i] * abs(i - j);
      G.addedge(i, j + n, 1, -gain);
    }
  }

  for (int i = 0; i < n; ++i) {
    // S から従業員 i へと、容量 2, コスト 0 の枝を張る
    G.addedge(2 * n, i, 1, 0);
  }

  for (int j = 0; j < n; ++j) {
    // 仕事 j から T へと、容量 1, コスト 0 の枝を張る
    G.addedge(j + n, 2 * n + 1, 1, 0);
  }

  // 最小費用流を求める
  COST res = MinCostFlow(G, 2 * n, 2 * n + 1, n);

  ll cost = 0;
  // 誰がどの仕事に割当てられたかを出力する
  for (int i = 0; i < n; ++i) {
    for (auto e : G[i]) {
      // 元々の容量 (e.icap) が 1 で、フローが流れて容量 (e.cap) が 0
      // になった部分が割り当てられたところ
      if (e.icap == 1 && e.cap == 0) {
        cost += a[i] * abs(i - e.to + n);
      }
    }
  }
  cout << cost << en;
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cout.tie(0);

  solve();
  // ll t;cin>>t;REP(i,t) solve();

  return 0;
}