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

#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstring>
#include <deque>
#include <fstream>
#include <functional>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <time.h>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
#define popcount __builtin_popcount
using namespace std;
typedef long long int ll;
typedef pair<ll, int> P;

int main() {
  int n;
  cin >> n;
  vector<P> va(n);
  ll a[2020];
  for (int i = 0; i < n; i++)
    cin >> a[i];
  for (int i = 0; i < n; i++)
    va[i] = P(a[i], i);
  sort(va.begin(), va.end(), greater<P>());
  ll dp[2020][2020];
  for (int i = 0; i <= n; i++)
    for (int j = 0; j <= n - i; j++)
      dp[i][j] = -1e18;
  dp[0][0] = 0;
  for (int i = 0; i <= n; i++) {
    for (int j = 0; j < n - i; j++) {
      ll a1 = va[i + j].first, k = va[i + j].second;
      dp[i + 1][j] = max(dp[i + 1][j], abs(i - k) * a1 + dp[i][j]);
      dp[i][j + 1] = max(dp[i][j + 1], abs(n - j - 1 - k) * a1 + dp[i][j]);
    }
  }
  ll ans = -1e18;
  for (int i = 0; i <= n; i++)
    ans = max(ans, dp[i][n - i]);
  cout << ans << endl;
  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <array>
#include <bitset>
#include <cassert>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstring>
#include <deque>
#include <fstream>
#include <functional>
#include <iostream>
#include <map>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <time.h>
#include <unordered_map>
#include <unordered_set>
#include <utility>
#include <vector>
#define popcount __builtin_popcount
using namespace std;
typedef long long int ll;
typedef pair<int, int> P;
template <typename T> struct hungarian { // n<=m
  const T inf = numeric_limits<T>::max();
  int n, m, max_match, root;
  T max_cost;
  vector<vector<T>> cost;
  vector<T> lx, ly, slack;
  vector<int> xy, yx, prev, slackx;
  vector<bool> s, t;
  void update_labels() {
    T delta = inf;
    for (int y = 0; y < m; y++)
      if (!t[y])
        delta = min(delta, slack[y]);
    for (int x = 0; x < n; x++)
      if (s[x])
        lx[x] -= delta;
    for (int y = 0; y < m; y++)
      if (t[y])
        ly[y] += delta;
    for (int y = 0; y < m; y++)
      if (!t[y])
        slack[y] -= delta;
  }
  void add_to_tree(int x, int prevx) {
    s[x] = true;
    prev[x] = prevx;
    for (int y = 0; y < m; y++) {
      if (lx[x] + ly[y] - cost[x][y] < slack[y]) {
        slack[y] = lx[x] + ly[y] - cost[x][y];
        slackx[y] = x;
      }
    }
  }
  void augment() {
    if (max_match == n)
      return;
    fill(s.begin(), s.end(), false);
    fill(t.begin(), t.end(), false);
    fill(prev.begin(), prev.end(), -1);
    queue<int> que;
    for (int x = 0; x < n; x++) {
      if (xy[x] == -1) {
        root = x;
        que.push(x);
        prev[x] = -2;
        s[x] = true;
        break;
      }
    }
    for (int y = 0; y < m; y++) {
      slack[y] = lx[root] + ly[y] - cost[root][y];
      slackx[y] = root;
    }
    int x, y;
    while (1) {
      while (!que.empty()) {
        x = que.front();
        que.pop();
        for (y = 0; y < m; y++) {
          if (cost[x][y] == lx[x] + ly[y] && !t[y]) {
            if (yx[y] == -1)
              break;
            t[y] = true;
            que.push(yx[y]);
            add_to_tree(yx[y], x);
          }
        }
        if (y < m)
          break;
      }
      if (y < m)
        break;
      update_labels();
      for (y = 0; y < m; y++) {
        if (!t[y] && slack[y] == 0) {
          if (yx[y] == -1) {
            x = slackx[y];
            break;
          } else {
            t[y] = true;
            que.push(yx[y]);
            add_to_tree(yx[y], slackx[y]);
          }
        }
      }
      if (y < m)
        break;
    }
    if (y < m) {
      max_match++;
      for (int cx = x, cy = y, ty; cx != -2; cx = prev[cx], cy = ty) {
        ty = xy[cx];
        yx[cy] = cx, xy[cx] = cy;
      }
      augment();
    }
  }
  hungarian(const vector<vector<T>> &cost)
      : max_match(0), max_cost(0), cost(cost), n(cost.size()),
        m(cost[0].size()), lx(n, -inf), ly(m), xy(n, -1), yx(m, -1), s(n), t(m),
        prev(n), slack(m), slackx(m) {
    for (int x = 0; x < n; x++)
      for (int y = 0; y < m; y++)
        lx[x] = max(lx[x], cost[x][y]);
    augment();
    for (int x = 0; x < n; x++)
      max_cost += cost[x][xy[x]];
  }
};
int main() {
  int n;
  cin >> n;
  ll a[2020];
  for (int i = 0; i < n; i++)
    cin >> a[i];
  vector<vector<ll>> c(n, vector<ll>(n));
  for (int i = 0; i < n; i++) {
    for (int j = 0; j < n; j++) {
      c[i][j] = a[i] * abs(i - j);
    }
  }
  hungarian<ll> h(c);
  cout << h.max_cost << endl;
  return 0;
}