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

#include <bits/stdc++.h>
using namespace std;
#define IOS                                                                    \\
  ios::sync_with_stdio(false);                                                 \\
  cin.tie(0);
#define FOR(i, s, n) for (int i = (s); i < (n); i++)
#define REP(i, n) FOR(i, 0, n)
#define RREP(i, n) for (int i = (n); i >= 0; i--)
#define ALL(n) (n).begin(), (n).end()
#define RALL(n) (n).rbegin(), (n).rend()
#define ATYN(n) cout << ((n) ? "Yes" : "No") << '\\n';
#define CFYN(n) cout << ((n) ? "YES" : "NO") << '\\n';
#define OUT(n) cout << (n) << '\\n';
using ll = long long;
using ull = unsigned long long;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
constexpr int MOD = 998244353;

int mod(const ll &x) { return x > 0 ? x % MOD : (x + MOD) % MOD; }

int modm(const ll &x, const ll &m) { return (x % MOD) * (m % MOD) % MOD; }

int modpow(ll x, ll p) {
  int ret = 1;
  while (p > 0) {
    if (p & 1)
      ret = modm(ret, x);
    x = modm(x, x);
    p >>= 1;
  }
  return ret;
}

// x * modinv(d)
int modd(ll x, ll d) { return modm(x, modpow(d, MOD - 2)); }

int main(void) {
  constexpr int FACT = 200000;
  int fact[FACT];
  fact[0] = 1;
  FOR(i, 1, FACT) fact[i] = modm(fact[i - 1], i);
  auto comb = [&](int x, int y) {
    int r = fact[x];
    r = modd(r, fact[y]);
    r = modd(r, fact[x - y]);
    return r;
  };
  // n種類のものから重複を許してr個配った時の組み合わせ数(重複組み合わせnHr)
  // n個のボールにr-1個仕切りを立てる、または、r個のボールにn-1個仕切りを立てる
  // auto multicomb = [&](int n,int r) {return comb(n+r-1,r);};

  IOS int N, M, K;
  cin >> N >> M >> K;
  int ans = 0;
  REP(i, K + 1) {
    int tmp = M;
    tmp = modm(tmp, comb(N - 1, i));
    tmp = modm(tmp, modpow(M - 1, N - 1 - i));
    ans = mod(ans + tmp);
  }

  cout << ans << '\\n';

  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define IOS                                                                    \\
  ios::sync_with_stdio(false);                                                 \\
  cin.tie(0);
#define FOR(i, s, n) for (int i = (s); i < (n); i++)
#define REP(i, n) FOR(i, 0, n)
#define RREP(i, n) for (int i = (n); i >= 0; i--)
#define ALL(n) (n).begin(), (n).end()
#define RALL(n) (n).rbegin(), (n).rend()
#define ATYN(n) cout << ((n) ? "Yes" : "No") << '\\n';
#define CFYN(n) cout << ((n) ? "YES" : "NO") << '\\n';
#define OUT(n) cout << (n) << '\\n';
using ll = long long;
using ull = unsigned long long;
using pii = pair<int, int>;
using pll = pair<ll, ll>;
constexpr int MOD = 998244353;

int mod(const ll &x) { return x > 0 ? x % MOD : (x + MOD) % MOD; }

int modm(const ll &x, const ll &m) { return (x % MOD) * (m % MOD) % MOD; }

int modpow(ll x, ll p) {
  int ret = 1;
  while (p > 0) {
    if (p & 1)
      ret = modm(ret, x);
    x = modm(x, x);
    p >>= 1;
  }
  return ret;
}

// x * modinv(d)
int modd(ll x, ll d) { return modm(x, modpow(d, MOD - 2)); }

int main(void) {
  IOS int N, M, K;
  cin >> N >> M >> K;
  int ans = modpow(M, N);
  cerr << ans << endl;

  vector<int> dp(K + 2);
  dp[0] = M;
  FOR(i, 1, N) {
    vector<int> tmp(K + 2);
    swap(tmp, dp);
    // cerr << max(1,K-N+i) << endl;
    FOR(j, max(1, K - N + i), K + 2) {
      dp[j] = mod(modm(tmp[j], M - 1) + tmp[j - 1]);
      // cerr << j << " " << dp[j] << endl;
    }
    dp[0] = modm(tmp[0], M - 1);
    int minus = modm(dp[K + 1], modpow(M, N - i - 1));
    dp[K + 1] = 0;
    // cerr << minus << endl;
    ans = mod(ans - minus);
  }

  cout << ans << '\\n';

  return 0;
}