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

/**
 * code generated by JHelper
 * More info: https://github.com/AlexeyDmitriev/JHelper
 * @author
 */

#include <fstream>
#include <iostream>

#include <bits/stdc++.h>

#define pb push_back
#define sz(v) ((int)(v).size())
#define all(v) (v).begin(), (v).end()
#define mp make_pair

using namespace std;

typedef long long int64;
typedef vector<int> vi;
typedef pair<int, int> ii;

const int64 MODULO = (int64)1e9 + 7;

typedef long long ll;
#define SZ 233333
const int MOD = 1e9 + 7; // or any prime
ll qp(ll a, ll b) {
  ll x = 1;
  a %= MOD;
  while (b) {
    if (b & 1)
      x = x * a % MOD;
    a = a * a % MOD;
    b >>= 1;
  }
  return x;
}
namespace linear_seq {
inline vector<int> BM(vector<int> x) {
  // ls: (shortest) relation sequence (after filling zeroes) so far
  // cur: current relation sequence
  vector<int> ls, cur;
  // lf: the position of ls (t')
  // ld: delta of ls (v')
  int lf, ld;
  for (int i = 0; i < int(x.size()); ++i) {
    ll t = 0;
    // evaluate at position i
    for (int j = 0; j < int(cur.size()); ++j)
      t = (t + x[i - j - 1] * (ll)cur[j]) % MOD;
    if ((t - x[i]) % MOD == 0)
      continue; // good so far
    // first non-zero position
    if (!cur.size()) {
      cur.resize(i + 1);
      lf = i;
      ld = (t - x[i]) % MOD;
      continue;
    }
    // cur=cur-c/ld*(x[i]-t)
    ll k = -(x[i] - t) * qp(ld, MOD - 2) % MOD /*1/ld*/;
    vector<int> c(i - lf - 1); // add zeroes in front
    c.pb(k);
    for (int j = 0; j < int(ls.size()); ++j)
      c.pb(-ls[j] * k % MOD);
    if (c.size() < cur.size())
      c.resize(cur.size());
    for (int j = 0; j < int(cur.size()); ++j)
      c[j] = (c[j] + cur[j]) % MOD;
    // if cur is better than ls, change ls to cur
    if (i - lf + (int)ls.size() >= (int)cur.size())
      ls = cur, lf = i, ld = (t - x[i]) % MOD;
    cur = c;
  }
  for (int i = 0; i < int(cur.size()); ++i)
    cur[i] = (cur[i] % MOD + MOD) % MOD;
  return cur;
}
int m; // length of recurrence
// a: first terms
// h: relation
ll a[SZ], h[SZ], t_[SZ], s[SZ], t[SZ];
// calculate p*q mod f
inline void mull(ll *p, ll *q) {
  for (int i = 0; i < m + m; ++i)
    t_[i] = 0;
  for (int i = 0; i < m; ++i)
    if (p[i])
      for (int j = 0; j < m; ++j)
        t_[i + j] = (t_[i + j] + p[i] * q[j]) % MOD;
  for (int i = m + m - 1; i >= m; --i)
    if (t_[i])
      // miuns t_[i]x^{i-m}(x^m-\\sum_{j=0}^{m-1} x^{m-j-1}h_j)
      for (int j = m - 1; ~j; --j)
        t_[i - j - 1] = (t_[i - j - 1] + t_[i] * h[j]) % MOD;
  for (int i = 0; i < m; ++i)
    p[i] = t_[i];
}
inline ll calc(ll K) {
  for (int i = m; ~i; --i)
    s[i] = t[i] = 0;
  // init
  s[0] = 1;
  if (m != 1)
    t[1] = 1;
  else
    t[0] = h[0];
  // binary-exponentiation
  while (K) {
    if (K & 1)
      mull(s, t);
    mull(t, t);
    K >>= 1;
  }
  ll su = 0;
  for (int i = 0; i < m; ++i)
    su = (su + s[i] * a[i]) % MOD;
  return (su % MOD + MOD) % MOD;
}
inline int work(vector<int> x, ll n) {
  if (n < int(x.size()))
    return x[n];
  vector<int> v = BM(x);
  m = v.size();
  if (!m)
    return 0;
  for (int i = 0; i < m; ++i)
    h[i] = v[i], a[i] = x[i];
  return calc(n);
}
} // namespace linear_seq

class TaskD {
public:
  void solveOne(istream &in, ostream &out) {
    int N, A;
    in >> N >> A;
    int bound = 3 * A;
    vector<vector<int>> w0(min(N, bound) + 1, vector<int>(A + 1));
    vector<vector<int>> w1(min(N, bound) + 1, vector<int>(A + 1));
    vector<int64> fact(N + 1);
    fact[0] = 1;
    for (int i = 1; i <= N; ++i) {
      fact[i] = i * fact[i - 1] % MODULO;
    }
    for (int n = 1; n < w0.size(); ++n) {
      for (int a = 1; a <= A && a <= n; ++a) {
        if (a == n) {
          w0[n][a] = fact[n];
          w1[n][a] = fact[n];
        } else {
          w0[n][a] = ((n - a) * (int64)w1[n - 1][a] +
                      (a - 1) * (int64)w1[n - 1][a - 1]) %
                     MODULO;
          w1[n][a] = (w0[n - 1][a - 1] + w0[n][a]) % MODULO;
        }
      }
    }
    vector<int> seq(w0.size() - A);
    for (int i = A; i < w0.size(); ++i)
      seq[i - A] = w0[i][A] * qp(fact[i - A], MODULO - 2) % MODULO;
    out << linear_seq::work(seq, N - A) * fact[N - A] % MODULO << "\\n";
  }

  void solve(std::istream &in, std::ostream &out) {
    int nt = 1;
    for (int it = 0; it < nt; ++it) {
      solveOne(in, out);
    }
  }
};

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  TaskD solver;
  std::istream &in(std::cin);
  std::ostream &out(std::cout);
  solver.solve(in, out);
  return 0;
}


Error: MLE

Fixed submission:

/**
 * code generated by JHelper
 * More info: https://github.com/AlexeyDmitriev/JHelper
 * @author
 */

#include <fstream>
#include <iostream>

#include <bits/stdc++.h>

#define pb push_back
#define sz(v) ((int)(v).size())
#define all(v) (v).begin(), (v).end()
#define mp make_pair

using namespace std;

typedef long long int64;
typedef vector<int> vi;
typedef pair<int, int> ii;

const int64 MODULO = (int64)1e9 + 7;

typedef long long ll;
#define SZ 233333
const int MOD = 1e9 + 7; // or any prime
ll qp(ll a, ll b) {
  ll x = 1;
  a %= MOD;
  while (b) {
    if (b & 1)
      x = x * a % MOD;
    a = a * a % MOD;
    b >>= 1;
  }
  return x;
}
namespace linear_seq {
inline vector<int> BM(vector<int> x) {
  // ls: (shortest) relation sequence (after filling zeroes) so far
  // cur: current relation sequence
  vector<int> ls, cur;
  // lf: the position of ls (t')
  // ld: delta of ls (v')
  int lf, ld;
  for (int i = 0; i < int(x.size()); ++i) {
    ll t = 0;
    // evaluate at position i
    for (int j = 0; j < int(cur.size()); ++j)
      t = (t + x[i - j - 1] * (ll)cur[j]) % MOD;
    if ((t - x[i]) % MOD == 0)
      continue; // good so far
    // first non-zero position
    if (!cur.size()) {
      cur.resize(i + 1);
      lf = i;
      ld = (t - x[i]) % MOD;
      continue;
    }
    // cur=cur-c/ld*(x[i]-t)
    ll k = -(x[i] - t) * qp(ld, MOD - 2) % MOD /*1/ld*/;
    vector<int> c(i - lf - 1); // add zeroes in front
    c.pb(k);
    for (int j = 0; j < int(ls.size()); ++j)
      c.pb(-ls[j] * k % MOD);
    if (c.size() < cur.size())
      c.resize(cur.size());
    for (int j = 0; j < int(cur.size()); ++j)
      c[j] = (c[j] + cur[j]) % MOD;
    // if cur is better than ls, change ls to cur
    if (i - lf + (int)ls.size() >= (int)cur.size())
      ls = cur, lf = i, ld = (t - x[i]) % MOD;
    cur = c;
  }
  for (int i = 0; i < int(cur.size()); ++i)
    cur[i] = (cur[i] % MOD + MOD) % MOD;
  return cur;
}
int m; // length of recurrence
// a: first terms
// h: relation
ll a[SZ], h[SZ], t_[SZ], s[SZ], t[SZ];
// calculate p*q mod f
inline void mull(ll *p, ll *q) {
  for (int i = 0; i < m + m; ++i)
    t_[i] = 0;
  for (int i = 0; i < m; ++i)
    if (p[i])
      for (int j = 0; j < m; ++j)
        t_[i + j] = (t_[i + j] + p[i] * q[j]) % MOD;
  for (int i = m + m - 1; i >= m; --i)
    if (t_[i])
      // miuns t_[i]x^{i-m}(x^m-\\sum_{j=0}^{m-1} x^{m-j-1}h_j)
      for (int j = m - 1; ~j; --j)
        t_[i - j - 1] = (t_[i - j - 1] + t_[i] * h[j]) % MOD;
  for (int i = 0; i < m; ++i)
    p[i] = t_[i];
}
inline ll calc(ll K) {
  for (int i = m; ~i; --i)
    s[i] = t[i] = 0;
  // init
  s[0] = 1;
  if (m != 1)
    t[1] = 1;
  else
    t[0] = h[0];
  // binary-exponentiation
  while (K) {
    if (K & 1)
      mull(s, t);
    mull(t, t);
    K >>= 1;
  }
  ll su = 0;
  for (int i = 0; i < m; ++i)
    su = (su + s[i] * a[i]) % MOD;
  return (su % MOD + MOD) % MOD;
}
inline int work(vector<int> x, ll n) {
  if (n < int(x.size()))
    return x[n];
  vector<int> v = BM(x);
  m = v.size();
  if (!m)
    return 0;
  for (int i = 0; i < m; ++i)
    h[i] = v[i], a[i] = x[i];
  return calc(n);
}
} // namespace linear_seq

class TaskD {
public:
  void solveOne(istream &in, ostream &out) {
    int N, A;
    in >> N >> A;
    int bound = 3 * A;
    vector<vector<int64>> w0(min(N, bound) + 1, vector<int64>(A + 1));
    vector<vector<int64>> w1(min(N, bound) + 1, vector<int64>(A + 1));
    vector<int64> fact(N + 1);
    fact[0] = 1;
    for (int i = 1; i <= N; ++i) {
      fact[i] = i * fact[i - 1] % MODULO;
    }
    for (int n = 1; n < w0.size(); ++n) {
      for (int a = 1; a <= A && a <= n; ++a) {
        if (a == n) {
          w0[n][a] = fact[n];
          w1[n][a] = fact[n];
        } else {
          w0[n][a] =
              ((n - a) * w1[n - 1][a] + (a - 1) * w1[n - 1][a - 1]) % MODULO;
          w1[n][a] = (w0[n - 1][a - 1] + w0[n][a]) % MODULO;
        }
      }
    }
    vector<int> seq(w0.size() - A);
    for (int i = A; i < w0.size(); ++i)
      seq[i - A] = w0[i][A] * qp(fact[i - A], MODULO - 2) % MODULO;
    out << linear_seq::work(seq, N - A) * fact[N - A] % MODULO << "\\n";
  }

  void solve(std::istream &in, std::ostream &out) {
    int nt = 1;
    for (int it = 0; it < nt; ++it) {
      solveOne(in, out);
    }
  }
};

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  TaskD solver;
  std::istream &in(std::cin);
  std::ostream &out(std::cout);
  solver.solve(in, out);
  return 0;
}