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

#include <algorithm>
#include <array>
#include <bitset>
#include <cmath>
#include <complex>
#include <cstdio>
#include <deque>
#include <fstream>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;

#define REP(i, n) for (int i = 0; i < n; ++i)
#define FOR(i, a, b) for (int i = a; i <= b; ++i)
#define FORR(i, a, b) for (int i = a; i >= b; --i)
#define ALL(c) (c).begin(), (c).end()

typedef long long ll;
typedef vector<int> VI;
typedef vector<ll> VL;
typedef vector<long double> VD;
typedef vector<VI> VVI;
typedef vector<VL> VVL;
typedef vector<VD> VVD;
typedef pair<int, int> P;
typedef pair<ll, ll> PL;
#define INF 1LL << 60

template <typename T> void chmin(T &a, T b) {
  if (a > b)
    a = b;
}
template <typename T> void chmax(T &a, T b) {
  if (a < b)
    a = b;
}

int in() {
  int x;
  scanf("%d", &x);
  return x;
}
ll lin() {
  ll x;
  scanf("%lld", &x);
  return x;
}

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

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

int main() {
  COMinit();
  ll N, M, K, mod = 998244353;
  cin >> N >> M >> K;

  ll ans = 0;
  vector<ll> color_combs(N);
  color_combs[0] = M;
  FOR(i, 1, N - 1) {
    color_combs[i] = color_combs[i - 1] * (M - 1);
    color_combs[i] %= mod;
  }

  REP(x, K + 1) {
    ll tmp = 1;
    tmp *= COM(N - 1, x);
    tmp *= color_combs[N - x - 1];
    tmp %= mod;
    ans += tmp;
    ans %= mod;
  }

  cout << ans << endl;
  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <array>
#include <bitset>
#include <cmath>
#include <complex>
#include <cstdio>
#include <deque>
#include <fstream>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>
using namespace std;

#define REP(i, n) for (int i = 0; i < n; ++i)
#define FOR(i, a, b) for (int i = a; i <= b; ++i)
#define FORR(i, a, b) for (int i = a; i >= b; --i)
#define ALL(c) (c).begin(), (c).end()

typedef long long ll;
typedef vector<int> VI;
typedef vector<ll> VL;
typedef vector<long double> VD;
typedef vector<VI> VVI;
typedef vector<VL> VVL;
typedef vector<VD> VVD;
typedef pair<int, int> P;
typedef pair<ll, ll> PL;
#define INF 1LL << 60

template <typename T> void chmin(T &a, T b) {
  if (a > b)
    a = b;
}
template <typename T> void chmax(T &a, T b) {
  if (a < b)
    a = b;
}

int in() {
  int x;
  scanf("%d", &x);
  return x;
}
ll lin() {
  ll x;
  scanf("%lld", &x);
  return x;
}

int main() {
  ll N, M, K, mod = 998244353;
  cin >> N >> M >> K;
  vector<vector<vector<ll>>> dp(
      N, vector<vector<ll>>(
             K + 1,
             vector<ll>(2))); // dp[i][j][k]:
                              // i番目のブロックまでみて、隣り合う組であって同じ色で塗られてるものがj組,
                              // k = 0ならば連続してない k = 1ならば連続してる
  dp[0][0][0] = M;
  FOR(i, 1, N - 1) {
    REP(j, min(ll(i + 1), K + 1)) {
      dp[i][j][0] =
          ((M - 1) * dp[i - 1][j][0] + (M - 1) * dp[i - 1][j][1]) % mod;
      if (j > 0) {
        dp[i][j][1] = (dp[i - 1][j - 1][0] + dp[i - 1][j - 1][1]) % mod;
      }
    }
    // REP(j, K+1) {
    //     cerr << "j = "<< j << ": " << dp[i][j][0] << " " << dp[i][j][1] << "
    //     ";
    // }
    // cerr << endl;
  }
  ll ans = 0;
  REP(j, K + 1) {
    ans += dp[N - 1][j][0] + dp[N - 1][j][1];
    ans %= mod;
  }
  cout << ans << endl;
  return 0;
}