Problem p00548 (C++) — buggy submission, verdict: Memory Limit Exceeded

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <utility>
#include <vector>

#define INF_LL 1e18
#define INF 1e9

#define REP(i, n) for (int i = 0; i < (n); i++)
#define FOR(i, a, b) for (int i = (a); i < (b); i++)
#define all(x) x.begin(), x.end()

#define fst first
#define snd second

using namespace std;

using ll = long long;
using PII = pair<int, int>;

class Union_find {
private:
  vector<int> par;
  vector<int> rank;
  int n;

public:
  Union_find(int a) {
    n = a;
    for (int i = 0; i < n; i++) {
      par.push_back(i);
      rank.push_back(0);
    }
  }

  int find(int x) {
    if (par[x] == x) {
      return x;
    } else {
      return par[x] = find(par[x]);
    }
  }

  void unite(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return;

    if (rank[x] < rank[y]) {
      par[x] = y;
    } else {
      par[y] = x;
      if (rank[x] == rank[y])
        rank[x]++;
    }
  }

  bool same(int x, int y) { return find(x) == find(y); }
};

pair<ll, pair<ll, ll>> dp[2][1010];
int main(void) {
  int N, M, K;
  ll A[20020];
  cin >> N >> M >> K;
  pair<ll, pair<ll, ll>> init = {INF_LL, {INF_LL, 0}};
  REP(i, 2) {
    REP(j, 1010) { dp[i][j] = init; }
  }
  REP(i, N) { cin >> A[i]; }
  dp[1][1] = {K, {A[0], A[0]}};
  FOR(i, 1, N) {
    REP(j, 1010) { dp[(i + 1) % 2][j] = init; }
    REP(j, min(i + 1, M + 1)) {
      if (0 < j && j < M) {
        ll cost = dp[i % 2][j].fst -
                  j * (dp[i % 2][j].snd.snd - dp[i % 2][j].snd.fst) +
                  (j + 1) * (max(dp[i % 2][j].snd.snd, A[i]) -
                             min(dp[i % 2][j].snd.fst, A[i]));
        if (dp[(i + 1) % 2][j + 1].fst > cost) {
          dp[(i + 1) % 2][j + 1] = {cost,
                                    {min(dp[i % 2][j].snd.fst, A[i]),
                                     max(dp[i % 2][j].snd.snd, A[i])}};
        }
      }
      ll cost = dp[i % 2][j].fst + K;
      if (dp[(i + 1) % 2][1].fst > cost) {
        dp[(i + 1) % 2][1] = {cost, {A[i], A[i]}};
      }
    }
  }
  ll res = INF_LL;
  REP(i, M + 1) { res = min(dp[N % 2][i].fst, res); }
  cout << res << endl;
}

Error: MLE

Fixed submission:

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <utility>
#include <vector>

#define INF_LL 1e18
#define INF 1e9

#define REP(i, n) for (int i = 0; i < (n); i++)
#define FOR(i, a, b) for (int i = (a); i < (b); i++)
#define all(x) x.begin(), x.end()

#define fst first
#define snd second

using namespace std;

using ll = long long;
using PII = pair<int, int>;

class Union_find {
private:
  vector<int> par;
  vector<int> rank;
  int n;

public:
  Union_find(int a) {
    n = a;
    for (int i = 0; i < n; i++) {
      par.push_back(i);
      rank.push_back(0);
    }
  }

  int find(int x) {
    if (par[x] == x) {
      return x;
    } else {
      return par[x] = find(par[x]);
    }
  }

  void unite(int x, int y) {
    x = find(x);
    y = find(y);
    if (x == y)
      return;

    if (rank[x] < rank[y]) {
      par[x] = y;
    } else {
      par[y] = x;
      if (rank[x] == rank[y])
        rank[x]++;
    }
  }

  bool same(int x, int y) { return find(x) == find(y); }
};

pair<ll, pair<ll, ll>> dp[20020][1010];
int main(void) {
  int N, M, K;
  ll A[20020];
  cin >> N >> M >> K;
  pair<ll, pair<ll, ll>> init = {INF_LL, {INF_LL, 0}};
  REP(i, 20020) {
    REP(j, 1010) { dp[i][j] = init; }
  }
  REP(i, N) { cin >> A[i]; }
  dp[1][1] = {K, {A[0], A[0]}};
  FOR(i, 1, N) {
    REP(j, min(i + 1, M + 1)) {
      if (0 < j && j < M) {
        ll cost = dp[i][j].fst - j * (dp[i][j].snd.snd - dp[i][j].snd.fst) +
                  (j + 1) * (max(dp[i][j].snd.snd, A[i]) -
                             min(dp[i][j].snd.fst, A[i]));
        if (dp[i + 1][j + 1].fst > cost) {
          dp[i + 1][j + 1] = {
              cost, {min(dp[i][j].snd.fst, A[i]), max(dp[i][j].snd.snd, A[i])}};
        }
      }
      ll cost = dp[i][j].fst + K;
      if (dp[i + 1][1].fst > cost) {
        dp[i + 1][1] = {cost, {A[i], A[i]}};
      }
    }
  }
  ll res = INF_LL;
  REP(i, M + 1) { res = min(dp[N][i].fst, res); }
  cout << res << endl;
}