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

#include <bits/stdc++.h>

// #define DEBUG

using namespace std;

const uint32_t MAX_K = 1000000;
uint64_t free_pow[MAX_K + 1];
uint64_t restrict_pow[MAX_K + 1];

const uint64_t MOD = 1000000007;

template <class T> void print_vect(std::vector<T> v) {
  for (int i = 0; i < v.size(); i++)
    cout << v[i] << " ";
  cout << endl;
}

uint64_t power_s(uint32_t a, uint32_t p, uint64_t M) {
  if (p < 3) {
    uint64_t ans = 1;
    while (p--) {
      ans = (ans * a) % M;
    }
    return ans;
  } else {
    uint64_t half_ans = power_s(a, p / 2, M);
    uint64_t ans = (half_ans * half_ans) % M;
    if (p % 2 == 1) {
      ans = (a * ans) % M;
    }
    return ans;
  }
}

void compute_pow() {
  uint64_t free = 1;
  uint64_t restrict = 1;
  free_pow[0] = free;
  restrict_pow[0] = restrict;
  for (uint32_t expo = 1; expo <= MAX_K; expo++) {
    free = free * 26;
    restrict = restrict * 25;
    if (free >= MOD) {
      free = free % MOD;
    }
    if (restrict >= MOD) {
      restrict = restrict % MOD;
    }
    free_pow[expo] = free;
    restrict_pow[expo] = restrict;
  }
}

void solve() {

  uint32_t K;
  cin >> K;
  std::string S;
  cin >> S;
  uint32_t last_position = S.size();
  K = K + S.size();
  uint64_t count = 0;

  compute_pow();
  uint64_t reposition_choices = 1;

  for (; last_position <= K; last_position++) {
    uint64_t last_choices = free_pow[K - last_position];
    uint64_t completion_choices = restrict_pow[last_position - S.size()];
    uint64_t overall_choices = last_choices * reposition_choices;
    if (overall_choices >= MOD) {
      overall_choices = overall_choices % MOD;
    }
    overall_choices = overall_choices * completion_choices;
    if (overall_choices >= MOD) {
      overall_choices = overall_choices % MOD;
    }
    count = (count + overall_choices) % MOD;

    reposition_choices = reposition_choices * last_position;
    if (reposition_choices >= MOD) {
      reposition_choices = reposition_choices % MOD;
    }

    uint64_t inverse_modulo =
        power_s(last_position + 1 - S.size(), MOD - 2, MOD);

    reposition_choices = reposition_choices * inverse_modulo;
    if (reposition_choices >= MOD) {
      reposition_choices = reposition_choices % MOD;
    }
  }
  cout << count << endl;
}

int main() {
  // making data IO Fast
  std::ios_base::sync_with_stdio(false);
  std::cin.tie(NULL);
  /****************************/

#ifdef DEBUG
  int TEST_CASES = 2;
  while (TEST_CASES--) {
#endif // DEBBUG

    solve();

#ifdef DEBUG
  }
#endif // DEBUG

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
const uint64_t MOD = 1000000007;

uint64_t power_s(uint32_t a, uint32_t p, uint64_t M) {
  if (p < 3) {
    uint64_t ans = 1;
    while (p--) {
      ans = (ans * a) % M;
    }
    return ans;
  } else {
    uint64_t half_ans = power_s(a, p / 2, M);
    uint64_t ans = (half_ans * half_ans) % M;
    if (p % 2 == 1) {
      ans = (a * ans) % M;
    }
    return ans;
  }
}

uint64_t combination(uint32_t m, uint32_t n, uint64_t M) {
  // choose m objects from n set object
  uint64_t ans = 1;
  for (uint32_t i = 1; i <= m; i++) {
    // assuming M prime
    uint64_t inverse_mod = power_s(i, M - 2, M);
    ans = (ans * (n + 1 - i)) % M;
    ans = (ans * inverse_mod) % M;
  }
  return ans;
}

void solve() {

  uint32_t K;
  cin >> K;
  std::string S;
  cin >> S;
  uint32_t last_position = S.size();
  K = K + S.size();
  uint64_t count = 0;

  for (; last_position <= K; last_position++) {
    uint64_t last_choices = power_s(26, K - last_position, MOD);
    uint64_t reposition_choices =
        combination(S.size() - 1, last_position - 1, MOD);
    uint64_t completion_choices = power_s(25, last_position - S.size(), MOD);
    uint64_t overall_choices = (last_choices * reposition_choices) % MOD;
    overall_choices = (overall_choices * completion_choices) % MOD;
    count = (count + overall_choices) % MOD;
  }
  cout << count << endl;
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);
  solve();
  return 0;
}