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

/*~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=
*$* WRITER:kakitamasziru/OxOmisosiru *$*
~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=*/
#ifdef LOCAL_JUDGE
#define _GLIBCXX_DEBUG // FOR THE DEBUG! COMMENT OUT THIS WHEN SUBMITTING!
#endif
/* I REALLY HOPE MY WISH REACH YOU , ATCODER'S ONLINE JUDGE */
#define WOULD
#define YOU
#define PLEASE
#define ACCEPT
#define MY
#define SUBMISSION
/* I REALLY HOPE MY WISH REACH YOU , ATCODER'S ONLINE JUDGE */
#include <algorithm> // min, max, swap, sort, reverse, lower_bound, upper_bound
#include <cstdint>   // int64_t, int*_t
#include <iomanip>
#include <iostream> // cout, endl, cin
#include <limits>   //setprecision
#include <string>   // string, to_string, stoi
#include <tuple>    // tuple, make_tuple
#include <utility>  // pair, make_pair
#include <vector>   // vector
// #include <cstdio> // printf
#include <bitset> // bitset
#include <cmath>  //abs,,,
#include <deque>  // deque
#include <map>    // map
#include <math.h> //pow,,,
#include <queue>  // queue, priority_queue
#include <set>    // set
#include <stack>  // stack
// It is so troublesome that I include bits/stdc++.h !
using namespace std;

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 long long INF = 100100100100;
const long long MOD = 998244353;
typedef pair<int, int> P;
// NのM乗を求める(繰り返し二乗法)
long long mod_pow(long long N, long long M) {
  if (M == 0)
    return 1;
  long long res = mod_pow((N * N) % MOD, M / 2);
  // 最下位ビット(*N)が1の時は単独でNをかける
  if (M & 1)
    res = (res * N) % MOD;

  return res %= MOD;
}
long long nCr(long long n, long long r) {
  long long A = 1, B = 1;
  // Aが分子Bが1/r!
  for (long long i = n - r + 1; i <= n; i++) {
    A = A * i % MOD;
  }
  for (long long i = 1; i <= r; i++) {
    B = B * i % MOD;
  }
  B = mod_pow(B, MOD - 2);
  B %= MOD;
  // Bは割るのではなく掛ける
  return (A * B) % MOD;
}

long long gcd(long long a, long long b) {
  if (b == 0)
    return a;
  else
    return gcd(b, a % b);
}

long long lcm(long long a, long long b) { return a * b / gcd(a, b); }

// long long A,B;
bool ok(long long key, long long index) {
  if (index > key)
    return true;
  else if (index <= key)
    return false;
}
long long binary_search(long long key, long long size) {
  // left,right,midはいずれもindex
  long long left = -1, right = size;
  while (right - left > 1) {
    long long mid = left + (right - left) / 2;
    long long hantei = ok(key, mid);
    if (hantei)
      right = mid;
    else
      left = mid;
  }
  if (left == -1)
    return 0;
  return left;
}

const int MAX = 510000;
// fac[]はn! inv[]はMODにおける逆元 finv[]はinv[]とfinv[]の累積積
long long fac[MAX], finv[MAX], inv[MAX];

// テーブルを作る前処理
void COMinit() {
  fac[0] = fac[1] = 1;
  finv[0] = finv[1] = 1;
  inv[1] = 1;
  for (int i = 2; i < MAX; i++) {
    fac[i] = fac[i - 1] * i % MOD;
    inv[i] = MOD - inv[MOD % i] * (MOD / i) % MOD;
    finv[i] = finv[i - 1] * inv[i] % MOD;
  }
}

// 二項係数計算(nCr)
long long COM(int n, int k) {
  if (n < k)
    return 0;
  if (n < 0 || k < 0)
    return 0;
  return fac[n] * (finv[k] * finv[n - k] % MOD) % MOD;
}

int main() {
  long long N, M, K;
  cin >> N >> M >> K;
  long long ans = 0;
  long long A = 1;
  COMinit();
  for (long long i = K; i >= 0; i--) {
    // 色の塗分けをグループごとで考えた時の場合の数
    // 累乗で求まるが愚直にやるとTLEするので、逆から見ていく(指数が大きくなるように見ていく)
    // ことで、毎回mod_powせずに*(N-i)できるようになる
    if (i == K)
      A = M * mod_pow(M - 1, N - i - 1);
    if (i != K)
      A *= M - 1;
    A %= MOD;
    // グループの分け方。N-1の間に接続部分を入れる場合の数
    long long B = COM(N - 1, i);
    long long sum = A * B;
    sum %= MOD;
    ans += sum;
    ans %= MOD;
    if (ans < 0)
      ans += MOD;
  }
  cout << ans % MOD << endl;
}


Error: TLE

Fixed submission:

/*~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=
*$* WRITER:kakitamasziru/OxOmisosiru *$*
~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=~=*/
#ifdef LOCAL_JUDGE
#define _GLIBCXX_DEBUG // FOR THE DEBUG! COMMENT OUT THIS WHEN SUBMITTING!
#endif
/* I REALLY HOPE MY WISH REACH YOU , ATCODER'S ONLINE JUDGE */
#define WOULD
#define YOU
#define PLEASE
#define ACCEPT
#define MY
#define SUBMISSION
/* I REALLY HOPE MY WISH REACH YOU , ATCODER'S ONLINE JUDGE */
#include <algorithm> // min, max, swap, sort, reverse, lower_bound, upper_bound
#include <cstdint>   // int64_t, int*_t
#include <iomanip>
#include <iostream> // cout, endl, cin
#include <limits>   //setprecision
#include <string>   // string, to_string, stoi
#include <tuple>    // tuple, make_tuple
#include <utility>  // pair, make_pair
#include <vector>   // vector
// #include <cstdio> // printf
#include <bitset> // bitset
#include <cmath>  //abs,,,
#include <deque>  // deque
#include <map>    // map
#include <math.h> //pow,,,
#include <queue>  // queue, priority_queue
#include <set>    // set
#include <stack>  // stack
// It is so troublesome that I include bits/stdc++.h !
using namespace std;

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 long long INF = 100100100100;
const long long MOD = 998244353;
typedef pair<int, int> P;
// NのM乗を求める(繰り返し二乗法)
long long mod_pow(long long N, long long M) {
  if (M == 0)
    return 1;
  long long res = mod_pow((N * N) % MOD, M / 2);
  // 最下位ビット(*N)が1の時は単独でNをかける
  if (M & 1)
    res = (res * N) % MOD;

  return res %= MOD;
}
long long nCr(long long n, long long r) {
  long long A = 1, B = 1;
  // Aが分子Bが1/r!
  for (long long i = n - r + 1; i <= n; i++) {
    A = A * i % MOD;
  }
  for (long long i = 1; i <= r; i++) {
    B = B * i % MOD;
  }
  B = mod_pow(B, MOD - 2);
  B %= MOD;
  // Bは割るのではなく掛ける
  return (A * B) % MOD;
}

long long gcd(long long a, long long b) {
  if (b == 0)
    return a;
  else
    return gcd(b, a % b);
}

long long lcm(long long a, long long b) { return a * b / gcd(a, b); }

// long long A,B;
bool ok(long long key, long long index) {
  if (index > key)
    return true;
  else if (index <= key)
    return false;
}
long long binary_search(long long key, long long size) {
  // left,right,midはいずれもindex
  long long left = -1, right = size;
  while (right - left > 1) {
    long long mid = left + (right - left) / 2;
    long long hantei = ok(key, mid);
    if (hantei)
      right = mid;
    else
      left = mid;
  }
  if (left == -1)
    return 0;
  return left;
}

int main() {
  long long N, M, K;
  cin >> N >> M >> K;
  long long ans = 0;
  long long A = 1;
  for (long long i = K; i >= 0; i--) {
    // 色の塗分けをグループごとで考えた時の場合の数
    // 累乗で求まるが愚直にやるとTLEするので、逆から見ていく(指数が大きくなるように見ていく)
    // ことで、毎回mod_powせずに*(N-i)できるようになる
    if (i == K)
      A = M * mod_pow(M - 1, N - i - 1);
    if (i != K)
      A *= M - 1;
    A %= MOD;
    // グループの分け方。N-1の間に接続部分を入れる場合の数
    long long B = nCr(N - 1, i);
    long long sum = A * B;
    sum %= MOD;
    ans += sum;
    ans %= MOD;
    if (ans < 0)
      ans += MOD;
  }
  cout << ans % MOD << endl;
}