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

#include <bits/stdc++.h>
#define REP(i, n) for (int i = 0; i < n; i++)
#define REPR(i, n) for (int i = n; i >= 0; i--)
#define FOR(i, m, n) for (int i = m; i < n; i++)
#define INF 2e9
#define MOD 1000000007
#define ALL(v) v.begin(), v.end()
using namespace std;
typedef long long ll;
using P = pair<int, int>;
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

const int mod = 1000000007;
struct mint {
  ll x; // typedef long long ll;
  mint(ll x = 0) : x((x % mod + mod) % mod) {}
  mint operator-() const { return mint(-x); }
  mint &operator+=(const mint a) {
    if ((x += a.x) >= mod)
      x -= mod;
    return *this;
  }
  mint &operator-=(const mint a) {
    if ((x += mod - a.x) >= mod)
      x -= mod;
    return *this;
  }
  mint &operator*=(const mint a) {
    (x *= a.x) %= mod;
    return *this;
  }
  mint operator+(const mint a) const { return mint(*this) += a; }
  mint operator-(const mint a) const { return mint(*this) -= a; }
  mint operator*(const mint a) const { return mint(*this) *= a; }
  mint pow(ll t) const {
    if (!t)
      return 1;
    mint a = pow(t >> 1);
    a *= a;
    if (t & 1)
      a *= *this;
    return a;
  }

  // for prime mod
  mint inv() const { return pow(mod - 2); }
  mint &operator/=(const mint a) { return *this *= a.inv(); }
  mint operator/(const mint a) const { return mint(*this) /= a; }
};
istream &operator>>(istream &is, const mint &a) { return is >> a.x; }
ostream &operator<<(ostream &os, const mint &a) { return os << a.x; }

vector<long long> fact, fact_inv, inv;

/*  init_nCk :二項係数のための前処理
    計算量:O(n)
*/
void init_nCk(int SIZE) {
  fact.resize(SIZE + 5);
  fact_inv.resize(SIZE + 5);
  inv.resize(SIZE + 5);
  fact[0] = fact[1] = 1;
  fact_inv[0] = fact_inv[1] = 1;
  inv[1] = 1;
  for (int i = 2; i < SIZE + 5; i++) {
    fact[i] = fact[i - 1] * i % MOD;
    inv[i] = MOD - inv[MOD % i] * (MOD / i) % MOD;
    fact_inv[i] = fact_inv[i - 1] * inv[i] % MOD;
  }
}

/*  nCk :MODでの二項係数を求める(前処理 int_nCk が必要)
    計算量:O(1)
*/
long long nCk(int n, int k) {
  assert(!(n < k));
  assert(!(n < 0 || k < 0));
  return fact[n] * (fact_inv[k] * fact_inv[n - k] % MOD) % MOD;
}

int main() {
  ll K;
  cin >> K;
  string S;
  cin >> S;

  mint ans2 = 0;
  vector<mint> vec(K + 1);
  vector<mint> vec2(K + 1);
  mint x = 1, y = 1;
  REP(i, K + 1) {
    vec.at(i) = x;
    x *= 25;
  }
  REP(i, K + 1) {
    vec2.at(i) = y;
    y *= 26;
  }
  init_nCk(S.size() + K);
  vector<ll> vec3(S.size() + K);
  FOR(i, S.size() - 1, S.size() + K) { vec3.at(i) = nCk(i, S.size() - 1); }
  FOR(i, S.size() - 1, S.size() + K) {
    mint ans = 1;
    int j = i - S.size() + 1;
    ans *= vec.at(j) * vec2.at(K - j);
    ans *= vec3.at(i);
    ans2 += ans;
  }

  cout << ans2 << endl;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define REP(i, n) for (int i = 0; i < n; i++)
#define REPR(i, n) for (int i = n; i >= 0; i--)
#define FOR(i, m, n) for (int i = m; i < n; i++)
#define INF 2e9
#define MOD 1000000007
#define ALL(v) v.begin(), v.end()
using namespace std;
typedef long long ll;
using P = pair<int, int>;
template <class T> inline bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

const int mod = 1000000007;
struct mint {
  ll x; // typedef long long ll;
  mint(ll x = 0) : x((x % mod + mod) % mod) {}
  mint operator-() const { return mint(-x); }
  mint &operator+=(const mint a) {
    if ((x += a.x) >= mod)
      x -= mod;
    return *this;
  }
  mint &operator-=(const mint a) {
    if ((x += mod - a.x) >= mod)
      x -= mod;
    return *this;
  }
  mint &operator*=(const mint a) {
    (x *= a.x) %= mod;
    return *this;
  }
  mint operator+(const mint a) const { return mint(*this) += a; }
  mint operator-(const mint a) const { return mint(*this) -= a; }
  mint operator*(const mint a) const { return mint(*this) *= a; }
  mint pow(ll t) const {
    if (!t)
      return 1;
    mint a = pow(t >> 1);
    a *= a;
    if (t & 1)
      a *= *this;
    return a;
  }

  // for prime mod
  mint inv() const { return pow(mod - 2); }
  mint &operator/=(const mint a) { return *this *= a.inv(); }
  mint operator/(const mint a) const { return mint(*this) /= a; }
};
istream &operator>>(istream &is, const mint &a) { return is >> a.x; }
ostream &operator<<(ostream &os, const mint &a) { return os << a.x; }
vector<long long> fact_inv, inv;

/*  init_nCk :二項係数のための前処理
    計算量:O(k)
*/
void init_nCk(int SIZE) {
  fact_inv.resize(SIZE + 5);
  inv.resize(SIZE + 5);
  fact_inv[0] = fact_inv[1] = 1;
  inv[1] = 1;
  for (int i = 2; i < SIZE + 5; i++) {
    inv[i] = MOD - inv[MOD % i] * (MOD / i) % MOD;
    fact_inv[i] = fact_inv[i - 1] * inv[i] % MOD;
  }
}

/*  nCk :MODでの二項係数を求める(前処理 int_nCk が必要)
    計算量:O(k)
*/
long long nCk(int n, int k) {
  assert(!(n < k));
  assert(!(n < 0 || k < 0));
  long long ans = 1;
  for (int i = n; i >= n - k + 1; i--) {
    ans *= i;
    ans %= MOD;
  }
  return ans * fact_inv[k] % MOD;
}

int main() {
  ll K;
  cin >> K;
  string S;
  cin >> S;

  mint ans2 = 0;
  vector<mint> vec(K + 1);
  vector<mint> vec2(K + 1);
  mint x = 1, y = 1;
  REP(i, K + 1) {
    vec.at(i) = x;
    x *= 25;
  }
  REP(i, K + 1) {
    vec2.at(i) = y;
    y *= 26;
  }
  init_nCk(S.size() - 1);
  vector<ll> vec3(S.size() + K);
  FOR(i, S.size() - 1, S.size() + K) { vec3.at(i) = nCk(i, S.size() - 1); }
  FOR(i, S.size() - 1, S.size() + K) {
    mint ans = 1;
    int j = i - S.size() + 1;
    ans *= vec.at(j) * vec2.at(K - j);
    ans *= vec3.at(i);
    ans2 += ans;
  }

  cout << ans2 << endl;
}