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

#include <cstdio>
#include <iostream>
#include <string>

using namespace std;

static const long long D = 998244353;

// 掛け算して剰余取る
long long Mul(long long a, long long b) {
  long long c = a * b;
  c %= D;
  return c;
}

long long *factMemo;
long long *modInvMemo;
long long *modInvMemoIntegrated;
long long *powMemo;

// nCk
long long Combination(int n, int k) {
  if (k == 0 || n == k)
    return 1;
  /*
  if (k == 0)
          return 1;
  long long a = 1;
  for (int i = n; i > n - k; i--) {
          a *= i;
          a %= D;
  }

  /*
  for (int i = 1; i <= k; i++) {
          a *= modInvMemo[i];
          a %= D;
  }
  */
  return Mul(Mul(factMemo[n], modInvMemoIntegrated[k]),
             modInvMemoIntegrated[n - k]);
  // return Mul(a, modInvMemoIntegrated[k]);
}

long long CalcModInv(long long a) {
  long long u = 1, v = 0;
  long long b = D;
  while (b) {
    long long t = a / b;
    a -= t * b;
    swap(a, b);
    u -= t * v;
    swap(u, v);
  }
  u %= D;
  if (u < 0)
    u += D;
  return u;
}

int main() {
  int N, M, K;
  {
    string s;
    getline(cin, s);
    sscanf(s.c_str(), "%d %d %d", &N, &M, &K);
  }

  factMemo = new long long[N];
  factMemo[0] = 1;
  for (int i = 1; i < N; i++) {
    factMemo[i] = Mul(factMemo[i - 1], i);
  }
  modInvMemo = new long long[N + 1];
  modInvMemoIntegrated = new long long[N + 1];
  modInvMemo[0] = 1;
  for (int i = 1; i <= N; i++) {
    modInvMemo[i] = CalcModInv(i);
  }
  // どうせ階乗でしか割らないので更にメモ
  modInvMemoIntegrated[1] = modInvMemo[1];
  modInvMemoIntegrated[0] = 1;
  for (int i = 2; i <= N; i++) {
    modInvMemoIntegrated[i] = Mul(modInvMemoIntegrated[i - 1], modInvMemo[i]);
  }

  powMemo = new long long[N + 1];
  powMemo[0] = 1;
  for (int i = 1; i <= N; i++) {
    powMemo[i] = Mul(powMemo[i - 1], M - 1);
  }

  // N 個
  // 最初の一個は無理
  // 以降の N - 1 個について、直前と同じ色になるものを k 個まで選ぶ
  //		sigma k = 0 to k = K
  // 逆に言えば 違う色になるものを K+1 個以上選ぶ ということになる
  // N - 1 - k 個について、直前と違う適当な色を塗る（M - 1 色ある）
  // 最初の 1 個はなんでもいいので M をかける

  // M * (N-1)C_k * (M-1)^(N-1-k)
  // これを sigma k = 0 to k = K
  // (N-1)C_k = (N-1)! / (k! * (N-1 - k)!)
  //

  long long ans = 0;
  for (int i = 0; i <= K; i++) {
    // long long a = 1;
    /*
    for (int j = 0; j < N - 1 - i; j++) {
            a = Mul(a, M - 1);
    }*/
    long long a = powMemo[N - 1 - i];
    a = Mul(a, Combination(N - 1, i));
    ans += a;
    ans %= D;
  }

  ans = Mul(ans, M);

  cout << ans << endl;

  delete[] factMemo;
  delete[] modInvMemo;
  delete[] modInvMemoIntegrated;

  return 0;
}

Error: TLE

Fixed submission:

#include <cstdio>
#include <iostream>
#include <string>
#pragma warning(disable : 4996)

using namespace std;

static const long long D = 998244353;

// 掛け算して剰余取る
long long Mul(long long a, long long b) {
  long long c = a * b;
  c %= D;
  return c;
}

long long *modInvMemo;
long long *modInvMemoIntegrated;
// nCk
long long Combination(int n, int k) {
  if (k == 0)
    return 1;
  long long a = 1;
  for (int i = n; i > n - k; i--) {
    a *= i;
    a %= D;
  }

  /*
  for (int i = 1; i <= k; i++) {
          a *= modInvMemo[i];
          a %= D;
  }
  */
  return Mul(a, modInvMemoIntegrated[k]);
}

long long CalcModInv(long long a) {
  long long u = 1, v = 0;
  long long b = D;
  while (b) {
    long long t = a / b;
    a -= t * b;
    swap(a, b);
    u -= t * v;
    swap(u, v);
  }
  u %= D;
  if (u < 0)
    u += D;
  return u;
}

int main() {
  int N, M, K;
  {
    string s;
    getline(cin, s);
    sscanf(s.c_str(), "%d %d %d", &N, &M, &K);
  }

  modInvMemo = new long long[K + 1];
  modInvMemoIntegrated = new long long[K + 1];
  for (int i = 1; i <= K; i++) {
    modInvMemo[i] = CalcModInv(i);
  }
  // どうせ階乗でしか割らないので更にメモ
  modInvMemoIntegrated[1] = modInvMemo[1];
  for (int i = 2; i <= K; i++) {
    modInvMemoIntegrated[i] = Mul(modInvMemoIntegrated[i - 1], modInvMemo[i]);
  }

  // N 個
  // 最初の一個は無理
  // 以降の N - 1 個について、直前と同じ色になるものを k 個まで選ぶ
  //		sigma k = 0 to k = K
  // 逆に言えば 違う色になるものを K+1 個以上選ぶ ということになる
  // N - 1 - k 個について、直前と違う適当な色を塗る（M - 1 色ある）
  // 最初の 1 個はなんでもいいので M をかける

  // M * (N-1)C_k * (M-1)^(N-1-k)
  // これを sigma k = 0 to k = K
  // (N-1)C_k = (N-1)! / (k! * (N-1 - k)!)
  //

  long long ans = 0;
  for (int i = 0; i <= K; i++) {
    long long a = 1;
    for (int j = 0; j < N - 1 - i; j++) {
      a = Mul(a, M - 1);
    }
    a = Mul(a, Combination(N - 1, i));
    ans += a;
    ans %= D;
  }

  ans = Mul(ans, M);

  cout << ans << endl;

  delete[] modInvMemo;

  return 0;
}