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

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

using i64 = int64_t;

#define rep(i, x, y)                                                           \\
  for (i64 i = i64(x), i##_max_for_repmacro = i64(y);                          \\
       i < i##_max_for_repmacro; ++i)
#define debug(x) #x << "=" << (x)

#ifdef DEBUG
#define _GLIBCXX_DEBUG
#define print(x) std::cerr << debug(x) << " (L:" << __LINE__ << ")" << std::endl
#else
#define print(x)
#endif

template <i64 p> class fp {
public:
  i64 x;
  fp() : x(0) {}
  fp(i64 x_) : x((x_ % p + p) % p) {}
  fp operator+() const { return fp(x); }
  fp operator-() const { return fp(-x); }
  fp &operator+=(const fp &y) {
    x += y.x;
    if (x >= p)
      x -= p;
    return *this;
  }
  fp &operator-=(const fp &y) { return *this += -y; }
  fp &operator*=(const fp &y) {
    x = x * y.x % p;
    return *this;
  }
  fp &operator/=(const fp &y) { return *this *= fp(inverse(y.x)); }
  fp operator+(const fp &y) const { return fp(x) += y; }
  fp operator-(const fp &y) const { return fp(x) -= y; }
  fp operator*(const fp &y) const { return fp(x) *= y; }
  fp operator/(const fp &y) const { return fp(x) /= y; }
  bool operator==(const fp &y) const { return x == y.x; }
  bool operator!=(const fp &y) const { return !(*this == y); }
  i64 extgcd(i64 a, i64 b, i64 &x, i64 &y) {
    i64 d = a;
    if (b != 0) {
      d = extgcd(b, a % b, y, x);
      y -= (a / b) * x;
    } else {
      x = 1;
      y = 0;
    }
    return d;
  }
  i64 inverse(i64 a) {
    i64 x, y;
    extgcd(a, p, x, y);
    return (x % p + p) % p;
  }
};

template <i64 p> i64 abs(const fp<p> &x) { return x.x; }

template <i64 p> istream &operator>>(istream &is, fp<p> &x) {
  is >> x.x;
  return is;
}

template <i64 p> ostream &operator<<(ostream &os, const fp<p> &x) {
  os << x.x;
  return os;
}

template <typename T, typename U>
ostream &operator<<(ostream &os, const pair<T, U> &p) {
  os << "(" << p.first << ", " << p.second << ")";
  return os;
}

template <typename T> ostream &operator<<(ostream &os, const vector<T> &vec) {
  os << "[";
  for (const auto &v : vec) {
    os << v << ",";
  }
  os << "]";
  return os;
}

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

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

template <typename A, typename T, size_t size>
void fill(A (&ary)[size], const T &val) {
  fill((T *)ary, (T *)(ary + size), val);
}

constexpr int inf = 1.01e9;
constexpr i64 inf64 = 4.01e18;
constexpr long double eps = 1e-9;

// double(64bit浮動小数)のn分探索のループ回数の上限(2分探索なら50でも十分かもしれない).
// long double(80ビットの x87 浮動小数点型?)だと,
// 2分探索であってもこれだと足りないケースがある気がするので,
// もうちょっと余裕を持たせた方が良さそう.
constexpr i64 max_loop = 100;

void solve() {
  // constexpr i64 mod = 1'000'000'007;

  i64 N;
  cin >> N;
  vector<i64> A(N);
  rep(i, 0, N) cin >> A[i];

  vector<pair<i64, i64>> ps(N);
  rep(i, 0, N) ps[i] = make_pair(A[i], i);

  sort(rbegin(ps), rend(ps));

  constexpr i64 maxn = 2001;
  static i64 memo[maxn][maxn];
  static bool done[maxn][maxn];
  fill(memo, -inf64);
  fill(done, false);

  function<i64(i64, i64)> rec = [&](i64 begin, i64 end) {
    auto &res = memo[begin][end];
    if (done[begin][end])
      return res;
    done[begin][end] = true;

    if (begin == end)
      return res = 0;

    i64 cnt_left = begin, cnt_right = N - end;
    i64 next_idx = cnt_left + cnt_right;

    chmax(res, rec(begin + 1, end) +
                   ps[next_idx].first * abs(ps[next_idx].second - begin));
    chmax(res, rec(begin, end - 1) +
                   ps[next_idx].first * abs(ps[next_idx].second - (end - 1)));

    return res;
  };

  cout << rec(0, N) << endl;
}

int main() {
  std::cin.tie(0);
  std::ios::sync_with_stdio(false);
  cout.setf(ios::fixed);
  cout.precision(16);
  solve();
  return 0;
}


Error: TLE

Fixed submission:

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

using i64 = int64_t;

#define rep(i, x, y)                                                           \\
  for (i64 i = i64(x), i##_max_for_repmacro = i64(y);                          \\
       i < i##_max_for_repmacro; ++i)
#define debug(x) #x << "=" << (x)

#ifdef DEBUG
#define _GLIBCXX_DEBUG
#define print(x) std::cerr << debug(x) << " (L:" << __LINE__ << ")" << std::endl
#else
#define print(x)
#endif

template <i64 p> class fp {
public:
  i64 x;
  fp() : x(0) {}
  fp(i64 x_) : x((x_ % p + p) % p) {}
  fp operator+() const { return fp(x); }
  fp operator-() const { return fp(-x); }
  fp &operator+=(const fp &y) {
    x += y.x;
    if (x >= p)
      x -= p;
    return *this;
  }
  fp &operator-=(const fp &y) { return *this += -y; }
  fp &operator*=(const fp &y) {
    x = x * y.x % p;
    return *this;
  }
  fp &operator/=(const fp &y) { return *this *= fp(inverse(y.x)); }
  fp operator+(const fp &y) const { return fp(x) += y; }
  fp operator-(const fp &y) const { return fp(x) -= y; }
  fp operator*(const fp &y) const { return fp(x) *= y; }
  fp operator/(const fp &y) const { return fp(x) /= y; }
  bool operator==(const fp &y) const { return x == y.x; }
  bool operator!=(const fp &y) const { return !(*this == y); }
  i64 extgcd(i64 a, i64 b, i64 &x, i64 &y) {
    i64 d = a;
    if (b != 0) {
      d = extgcd(b, a % b, y, x);
      y -= (a / b) * x;
    } else {
      x = 1;
      y = 0;
    }
    return d;
  }
  i64 inverse(i64 a) {
    i64 x, y;
    extgcd(a, p, x, y);
    return (x % p + p) % p;
  }
};

template <i64 p> i64 abs(const fp<p> &x) { return x.x; }

template <i64 p> istream &operator>>(istream &is, fp<p> &x) {
  is >> x.x;
  return is;
}

template <i64 p> ostream &operator<<(ostream &os, const fp<p> &x) {
  os << x.x;
  return os;
}

template <typename T, typename U>
ostream &operator<<(ostream &os, const pair<T, U> &p) {
  os << "(" << p.first << ", " << p.second << ")";
  return os;
}

template <typename T> ostream &operator<<(ostream &os, const vector<T> &vec) {
  os << "[";
  for (const auto &v : vec) {
    os << v << ",";
  }
  os << "]";
  return os;
}

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

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

template <typename A, typename T, size_t size>
void fill(A (&ary)[size], const T &val) {
  fill((T *)ary, (T *)(ary + size), val);
}

constexpr int inf = 1.01e9;
constexpr i64 inf64 = 4.01e18;
constexpr long double eps = 1e-9;

// double(64bit浮動小数)のn分探索のループ回数の上限(2分探索なら50でも十分かもしれない).
// long double(80ビットの x87 浮動小数点型?)だと,
// 2分探索であってもこれだと足りないケースがある気がするので,
// もうちょっと余裕を持たせた方が良さそう.
constexpr i64 max_loop = 100;

struct MinimumCostMaximumFlow {
  typedef i64 Index;
  typedef i64 Flow;
  typedef i64 Cost;
  static const Flow InfCapacity = inf64;
  static const Cost InfCost = inf64;
  struct Edge {
    Index to;
    Index rev;
    Flow capacity;
    Cost cost;
  };
  vector<vector<Edge>> g;
  void init(Index n) { g.assign(n, vector<Edge>()); }
  void add(Index i, Index j, Flow capacity = InfCapacity, Cost cost = Cost()) {
    Edge e, f;
    e.to = j, f.to = i;
    e.capacity = capacity, f.capacity = 0;
    e.cost = cost, f.cost = -cost;
    g[i].push_back(e);
    g[j].push_back(f);
    g[i].back().rev = (Index)g[j].size() - 1;
    g[j].back().rev = (Index)g[i].size() - 1;
  }
  void addB(Index i, Index j, Flow capacity = InfCapacity, Cost cost = Cost()) {
    add(i, j, capacity, cost);
    add(j, i, capacity, cost);
  }
  pair<Cost, Flow> minimumCostMaximumFlow(Index s, Index t,
                                          Flow f = InfCapacity,
                                          bool useSPFA = false) {
    int n = g.size();
    vector<Cost> dist(n);
    vector<Index> prev(n);
    vector<Index> prevEdge(n);
    pair<Cost, Flow> total = make_pair(0, 0);
    vector<Cost> potential(n);
    while (f > 0) {
      fill(dist.begin(), dist.end(), InfCost);
      if (useSPFA || total.second == 0) {
        deque<Index> q;
        q.push_back(s);
        dist[s] = 0;
        vector<bool> inqueue(n);
        while (!q.empty()) {
          Index i = q.front();
          q.pop_front();
          inqueue[i] = false;
          for (Index ei = 0; ei < g[i].size(); ei++) {
            const Edge &e = g[i][ei];
            Index j = e.to;
            Cost d = dist[i] + e.cost;
            if (e.capacity > 0 && d < dist[j]) {
              if (!inqueue[j]) {
                inqueue[j] = true;
                q.push_back(j);
              }
              dist[j] = d;
              prev[j] = i;
              prevEdge[j] = ei;
            }
          }
        }
      } else {
        vector<bool> vis(n);
        priority_queue<pair<Cost, Index>> q;
        q.push(make_pair(-0, s));
        dist[s] = 0;
        while (!q.empty()) {
          Index i = q.top().second;
          q.pop();
          if (vis[i])
            continue;
          vis[i] = true;
          for (Index ei = 0; ei < (Index)g[i].size(); ei++) {
            const Edge &e = g[i][ei];
            if (e.capacity <= 0)
              continue;
            Index j = e.to;
            Cost d = dist[i] + e.cost + potential[i] - potential[j];
            if (dist[j] > d) {
              dist[j] = d;
              prev[j] = i;
              prevEdge[j] = ei;
              q.push(make_pair(-d, j));
            }
          }
        }
      }
      if (dist[t] == InfCost)
        break;
      if (!useSPFA)
        for (Index i = 0; i < n; i++)
          potential[i] += dist[i];

      Flow d = f;
      Cost distt = 0;
      for (Index v = t; v != s;) {
        Index u = prev[v];
        const Edge &e = g[u][prevEdge[v]];
        d = min(d, e.capacity);
        distt += e.cost;
        v = u;
      }
      f -= d;
      total.first += d * distt;
      total.second += d;
      for (Index v = t; v != s; v = prev[v]) {
        Edge &e = g[prev[v]][prevEdge[v]];
        e.capacity -= d;
        g[e.to][e.rev].capacity += d;
      }
    }
    return total;
  }
};
void solve() {
  // constexpr i64 mod = 1'000'000'007;

  i64 N;
  cin >> N;
  vector<i64> A(N);
  rep(i, 0, N) cin >> A[i];

  const i64 s = 2 * N, t = 2 * N + 1;
  MinimumCostMaximumFlow mcf;
  mcf.init(2 * N + 2);
  rep(i, 0, N) {
    mcf.add(s, i, 1, 0);
    mcf.add(i + N, t, 1, 0);
  }
  rep(i, 0, N) {
    rep(j, 0, N) { mcf.add(i, j + N, 1, -A[i] * abs(i - j)); }
  }

  cout << -mcf.minimumCostMaximumFlow(s, t, N, true).first << endl;
}

int main() {
  std::cin.tie(0);
  std::ios::sync_with_stdio(false);
  cout.setf(ios::fixed);
  cout.precision(16);
  solve();
  return 0;
}