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

// #define _GLIBCXX_DEBUG // for STL debug (optional)
#include <algorithm>
#include <bitset>
#include <cassert>
#include <cfloat>
#include <climits>
#include <cmath>
#include <complex>
#include <cstdio>
#include <cstring>
#include <ctime>
#include <deque>
#include <fstream>
#include <functional>
#include <iomanip>
#include <iostream>
#include <limits>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <utility>
#include <vector>
using namespace std;
using ll = long long int;
using int64 = long long int;

template <typename T> void chmax(T &a, T b) { a = max(a, b); }
template <typename T> void chmin(T &a, T b) { a = min(a, b); }
template <typename T> void chadd(T &a, T b) { a = a + b; }

int dx[] = {0, 0, 1, -1};
int dy[] = {1, -1, 0, 0};
const int INF = 1LL << 29;
const ll LONGINF = 1LL << 60;
const ll MOD = 998244353LL;

// ModInt begin

using ll = long long;
template <ll mod> struct ModInt {
  ll v;
  ll mod_pow(ll x, ll n) const {
    return (!n) ? 1 : (mod_pow((x * x) % mod, n / 2) * ((n & 1) ? x : 1)) % mod;
  }
  ModInt(ll a = 0) : v((a %= mod) < 0 ? a + mod : a) {}
  ModInt operator+(const ModInt &b) const {
    return (v + b.v >= mod ? ModInt(v + b.v - mod) : ModInt(v + b.v));
  }
  ModInt operator-() const { return ModInt(-v); }
  ModInt operator-(const ModInt &b) const {
    return (v - b.v < 0 ? ModInt(v - b.v + mod) : ModInt(v - b.v));
  }
  ModInt operator*(const ModInt &b) const { return (v * b.v) % mod; }
  ModInt operator/(const ModInt &b) const {
    return (v * mod_pow(b.v, mod - 2)) % mod;
  }

  bool operator==(const ModInt &b) const { return v == b.v; }
  bool operator!=(const ModInt &b) const { return !(*this == b); }
  ModInt &operator+=(const ModInt &b) {
    v += b.v;
    if (v >= mod)
      v -= mod;
    return *this;
  }
  ModInt &operator-=(const ModInt &b) {
    v -= b.v;
    if (v < 0)
      v += mod;
    return *this;
  }
  ModInt &operator*=(const ModInt &b) {
    (v *= b.v) %= mod;
    return *this;
  }
  ModInt &operator/=(const ModInt &b) {
    (v *= mod_pow(b.v, mod - 2)) %= mod;
    return *this;
  }
  ModInt pow(ll x) { return ModInt(mod_pow(v, x)); }
  // operator int() const { return int(v); }
  // operator long long int() const { return v; }
};

template <ll mod> ModInt<mod> pow(ModInt<mod> n, ll k) {
  return ModInt<mod>(n.mod_pow(n.v, k));
}

template <ll mod> ostream &operator<<(ostream &out, ModInt<mod> a) {
  return out << a.v;
}
template <ll mod> istream &operator>>(istream &in, ModInt<mod> &a) {
  in >> a.v;
  return in;
}

// ModInt end

// 各種組み合わせを求めるライブラリ
template <typename NumType> struct Combination {
  int LIMIT;
  vector<NumType> fact_, finv_;

  Combination() {}
  Combination(int LIMIT_) : LIMIT(LIMIT_), fact_(LIMIT + 1), finv_(LIMIT + 1) {
    fact_[0] = finv_[LIMIT] = NumType(1);
    for (int i = 1; i <= LIMIT; i++) {
      fact_[i] = fact_[i - 1] * NumType(i);
    }

    finv_[LIMIT] /= fact_[LIMIT];
    for (int i = LIMIT - 1; i >= 0; i--) {
      finv_[i] = finv_[i + 1] * NumType(i + 1);
    }
  }

  inline NumType fact(int k) const { return fact_[k]; }
  inline NumType finv(int k) const { return finv_[k]; }
  NumType P(int n, int r) const {
    if (r < 0 or n < r)
      return NumType(0);
    return fact_[n] * finv_[n - r];
  }
  NumType C(int n, int r) const {
    if (r < 0 or n < r)
      return NumType(0);
    return fact_[n] * finv_[n - r] * finv_[r];
  }
  // 重複組み合わせ
  NumType H(int n, int r) const {
    if (n < 0 or r < 0)
      return NumType(0);
    return r == 0 ? NumType(1) : C(n + r - 1, r);
  }
  // ベル数 (区別できる n 個のボールを区別できない k 個以下の箱に分割)
  // B(n, n) := n 個のボールを任意個のグループに分割する場合の数
  NumType B(int n, int k) const {
    if (n == 0)
      return NumType(1);
    k = min(n, k);
    NumType ret(0);
    vector<NumType> pref(k + 1);
    pref[0] = NumType(1);
    for (int i = 1; i <= k; i++) {
      if (i & 1)
        pref[i] = pref[i - 1] - finv_[i];
      else
        pref[i] = pref[i - 1] + finv_[i];
    }
    for (int i = 1; i <= k; i++) {
      // 累乗が必要なので適宜書き換える？
      // ModInt 使うならこれでいい
      ret += NumType(i).pow(n) * finv_[i] * pref[k - i];
    }
    return ret;
  }
  // スターリング数 (区別できる n 個のボールを区別できない k 個の箱に分割)
  NumType S(int n, int k) const {
    if (n < k)
      return NumType(0);
    NumType ans(0);
    for (int i = 0; i <= k; i++) {
      NumType val = C(k, i) * NumType(i).pow(n);
      if ((k - i) % 2)
        ans -= val;
      else
        ans += val;
    }
    return ans * finv_[k];
  }

  // ランダムウォーク: 左に X 回、右に Y 回進むとき、
  // 移動途中・終了時に座標 K を超えないものを数える
  // K が非負なら左側、負なら右側領域 (？) のランダムウォーク
  NumType walk(int X, int Y, int K) {
    if (K < 0)
      return walk(Y, X, -K);
    if (Y <= K)
      return C(X + Y, X); // 引っかからない
    if (Y - X > K)
      return NumType(0); // そもそも合計が超える

    int A = Y - K - 1, B = X + K + 1;
    return C(X + Y, X) - C(A + B, A);
  }
};

// P(n, k) := n の k 分割 (k 個の 0 以上の整数の和)
template <typename NumType, int LIMIT = 3010> struct Partition {
  vector<vector<NumType>> dp;
  Partition() : dp(LIMIT, vector<NumType>(LIMIT)) {
    for (int k = 0; k < LIMIT; k++)
      dp[0][k] = NumType(1);
    for (int i = 1; i < LIMIT; i++) {
      for (int j = 1; j < LIMIT; j++) {
        dp[i][j] += dp[i][j - 1];
        if (i - j >= 0)
          dp[i][j] += dp[i - j][j];
      }
    }
  }
  inline NumType get(int n, int k) {
    if (n < 0 or k < 0)
      return NumType(0);
    return dp[n][k];
  }
};

using mint = ModInt<MOD>;

int main() {
  Combination<mint> comb(500010);
  ll N, M, K;
  cin >> N >> M >> K;
  mint ans(0);
  for (int x = 0; x <= K; x++) {
    mint mul = mint(M) * pow(mint(M - 1), N - x - 1) * comb.H(N - x, x);
    ans += mul;
  }
  cout << ans << endl;
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;

const int MOD = 998244353;
const int MAX = 510000;

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;
  }
}

// 二項係数計算
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;
}

long long modpow(long long x, long long n, long long mod) {
  long long res = 1;
  for (int i = 0; i < n; i++) {
    res *= x;
    res %= mod;
  }
  return res;
}

int main() {

  long long N, M, K;
  cin >> N >> M >> K;

  // 前処理
  COMinit();

  // 計算例
  long long sum = 0;
  for (int k = 0; k <= K; k++) {
    sum += ((M * COM(N - 1, k) % MOD) * modpow(M - 1, N - 1 - k, MOD)) % MOD;
    sum %= MOD;
    // cout << modpow(M-1,N-1-k,MOD) << endl;
  }
  cout << sum << endl;
}