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


#pragma GCC optimize("O3")

#include <bits/stdc++.h>

using namespace std;

#ifdef LOCAL
#include <prettyprint.hpp>
#define debug(...) cerr << "[" << #__VA_ARGS__ << "]: ", d_err(__VA_ARGS__);
#else
#define debug(...) 83;
#endif

void d_err() { cerr << endl; }

template <typename H, typename... T> void d_err(H h, T... t) {
  cerr << h << " ";
  d_err(t...);
}

template <typename T> void print(T x) { cout << x << endl; }

#define ALL(x) (x).begin(), (x).end()
#define FOR(i, m, n) for (int i = (m); i < (n); ++i)
#define REVFOR(i, m, n) for (int i = (n - 1); i >= (m); --i)
#define REP(i, n) FOR(i, 0, n)
#define REVREP(i, n) REVFOR(i, 0, n)
#define fi first
#define se second
#define pb push_back
#define mp make_pair
#define eb emplace_back
#define bcnt __builtin_popcountll

typedef long long ll;
typedef vector<int> vi;
typedef vector<ll> vll;
typedef pair<ll, ll> Pll;
typedef pair<int, int> Pin;

ll INF = 1e16;
int inf = 1e9;
ll MOD = 1e9 + 7;

int main() {
  cin.tie(0);
  ios_base::sync_with_stdio(false);
  cout << fixed << setprecision(20);

  int N;
  cin >> N;

  vector<pair<ll, int>> A(N);
  REP(i, N) {
    ll a;
    cin >> a;
    A[i] = {a, i};
  }

  sort(ALL(A));
  reverse(ALL(A));

  ll dp[N + 2][N + 2];
  memset(dp, 0, sizeof(dp));

  REP(i, N) {
    REP(y, i + 2) {
      int x = i - y + 1;
      if (x > 0)
        dp[x][y] = max(dp[x][y], dp[x - 1][y] + A[i].fi * (A[i].se - x + 1));
      if (y > 0)
        dp[x][y] = max(dp[x][y], dp[x][y - 1] + A[i].fi * (N - y - A[i].se));
    }
  }

  ll ans = 0;
  REP(i, N) { ans = max(ans, dp[N - i][i]); }
  print(ans);
}


Error: TLE

Fixed submission:


#pragma GCC optimize("O3")

#include <bits/stdc++.h>

using namespace std;

#ifdef LOCAL
#include <prettyprint.hpp>
#define debug(...) cerr << "[" << #__VA_ARGS__ << "]: ", d_err(__VA_ARGS__);
#else
#define debug(...) 83;
#endif

void d_err() { cerr << endl; }

template <typename H, typename... T> void d_err(H h, T... t) {
  cerr << h << " ";
  d_err(t...);
}

template <typename T> void print(T x) { cout << x << endl; }

#define ALL(x) (x).begin(), (x).end()
#define FOR(i, m, n) for (int i = (m); i < (n); ++i)
#define REVFOR(i, m, n) for (int i = (n - 1); i >= (m); --i)
#define REP(i, n) FOR(i, 0, n)
#define REVREP(i, n) REVFOR(i, 0, n)
#define fi first
#define se second
#define pb push_back
#define mp make_pair
#define eb emplace_back
#define bcnt __builtin_popcountll

typedef long long ll;
typedef vector<int> vi;
typedef vector<ll> vll;
typedef pair<ll, ll> Pll;
typedef pair<int, int> Pin;

ll INF = 1e16;
int inf = 1e9;
ll MOD = 1e9 + 7;

const int MAX_V = 1e4;
const int MAX_E = 1e4;

struct Edge {
  int from, to, w;
};

vector<Edge> w(MAX_E);
ll d[MAX_V];
int V, E;

template <typename T> struct PrimalDual {
  struct Graph {
    int to, rev;
    T cap;
    T cost;

    Graph(int t, int r, const T &cp, const T &cs)
        : to(t), rev(r), cap(cp), cost(cs){};
  };

  vector<T> d;
  vector<vector<Graph>> g;
  vector<int> prevv, preve;
  vector<bool> used;
  int sz;

  PrimalDual(int N) : sz(N) {
    d.resize(N);
    g.resize(N);
    prevv.resize(N);
    preve.resize(N);
    used.resize(N);
  }

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

  void bellman_ford(int s) {
    d.assign(sz, INF);
    d[s] = 0;

    bool update = true;

    while (update) {
      update = false;
      REP(i, sz) {
        if (d[i] == INF)
          continue;
        REP(j, g[i].size()) {
          auto &e = g[i][j];
          if (e.cap > 0 && d[e.to] > d[i] + e.cost) {
            d[e.to] = d[i] + e.cost;
            prevv[e.to] = i;
            preve[e.to] = j;
            update = true;
          }
        }
      }
    }
  }

  T operator()(int start, int end, T f) {
    T res = 0;
    while (f > 0) {
      bellman_ford(start);

      if (d[end] == INF)
        return -1;

      T f_cur = f;
      for (int v = end; v != start; v = prevv[v]) {
        f_cur = min(f, g[prevv[v]][preve[v]].cap);
      }
      f -= f_cur;

      res += d[end] * f_cur;

      for (int v = end; v != start; v = prevv[v]) {
        auto &e = g[prevv[v]][preve[v]];
        e.cap -= f_cur;
        g[v][e.rev].cap += f_cur;
      }
    }
    return res;
  }
};

int main() {
  cin.tie(0);
  ios_base::sync_with_stdio(false);
  cout << fixed << setprecision(20);

  int N;
  cin >> N;

  PrimalDual<ll> pd(2 * N + 2);

  REP(i, N) {
    ll a;
    cin >> a;
    REP(j, N) { pd.add_edge(i, N + j, 1, -a * abs(i - j)); }

    pd.add_edge(2 * N, i, 1, 0);
    pd.add_edge(N + i, 2 * N + 1, 1, 0);
  }

  print(-pd(2 * N, 2 * N + 1, N));
}