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

#include <bits/stdc++.h>

using namespace std;
mt19937_64 rng(58);

typedef long long ll;

int k;
string s;
// find b s.t. a * b = 1 (mod m)
template <typename T> T inverse(T a, T m) {
  T u = 0, v = 1;
  while (a != 0) {
    T t = m / a;
    m -= t * a;
    swap(a, m);
    u -= t * v;
    swap(u, v);
  }
  assert(m == 1);
  return u;
}

template <long long mod = 1000000007> struct modint {
  long long a;
  modint() : a(0) {}
  modint(long long t) {
    a = t % mod;
    if (a < 0)
      a += mod;
  }
  operator long long() const { return a; }
  bool operator==(const modint &x) const { return a == x.a; }
  bool operator!=(const modint &x) const { return a != x.a; }
  modint operator-() const { return modint(a ? (mod - a) : 0); }
  modint operator~() const { return inverse(a, mod); }

  modint operator+(const modint &x) const { return modint(*this) += x; }
  modint operator-(const modint &x) const { return modint(*this) -= x; }
  modint operator*(const modint &x) const { return modint(*this) *= x; }
  modint operator/(const modint &x) const { return modint(*this) /= x; }

  modint &operator+=(const modint &x) {
    a += x.a;
    if (a >= mod)
      a -= mod;
    return *this;
  }
  modint &operator-=(const modint &x) {
    a -= x.a;
    if (a < 0)
      a += mod;
    return *this;
  }
  modint &operator*=(const modint &x) {
    a = a * x.a % mod;
    return *this;
  }
  modint &operator/=(const modint &x) {
    a = a * (~x).a % mod;
    return *this;
  }

  friend ostream &operator<<(ostream &os, const modint &x) { return os << x.a; }
  friend istream &operator>>(istream &is, modint &x) {
    long long t;
    is >> t;
    x = modint(t);
    return is;
  }

  modint pow(long long x) const {
    modint ret = 1, tmp = a;
    for (; x; tmp *= tmp, x >>= 1) {
      if (x & 1)
        ret *= tmp;
    }
    return ret;
  }
};

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cout.tie(NULL);
  cin >> k >> s;

  modint<> f[k + s.length() + 1], finv[k + s.length() + 1];
  f[0] = finv[0] = 1;
  for (int i = 1; i <= k + s.length(); i++) {
    f[i] = i * f[i - 1];
    finv[i] = ~f[i];
  }

  modint<> ans = 0;
  for (int idx = s.length() - 1; idx < k + s.length(); idx++) {
    ans += f[idx] * finv[s.length() - 1] * finv[idx - s.length() + 1] *
           modint<>(25).pow(idx - s.length() + 1) *
           modint<>(26).pow(s.length() + k - 1 - idx);
  }
  cout << ans << endl;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;
mt19937_64 rng(58);

typedef long long ll;
const ll mod = 1000000007;

int k;
string s;

ll power(ll x, int p) {
  ll r = 1;
  while (p) {
    if (p % 2)
      r = (r * x) % mod;
    x = (x * x) % mod;
    p >>= 1;
  }
  return r % mod;
}

map<pair<int, int>, ll> mem;

ll f(int k, int s) {
  if (k < 0)
    return 0;
  if (k == 0)
    return 1;
  if (s == 0)
    return power(26, k);
  if (mem.find({k, s}) != mem.end())
    return mem[{k, s}];

  ll ans = 0;
  for (int i = 0; i <= k; i++) {
    ans = (ans + power(25, i) * f(k - i, s - 1)) % mod;
  }
  return mem[{k, s}] = ans;
}

const int nax = 1000000 + 1;
ll dp[nax];

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  cout.tie(NULL);
  cin >> k >> s;

  for (int i = 0; i <= k; i++)
    dp[i] = power(26, i);
  for (int t = 1; t <= s.length(); t++) {
    for (int i = 1; i <= k; i++) {
      dp[i] = (dp[i] + 25 * dp[i - 1]) % mod;
    }
  }
  // cout << f(k, s.length()) << endl;
  cout << dp[k] << endl;
}