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

#include <bits/stdc++.h>
#define INF 1000000007
#define MOD 998244353
#define rep(i, N) for (ll i = 0; i < N; i++)
#define rep2(i, j, k) for (ll i = j; i < k; i++)
#define chmax(x, y) x = max(x, y)
#define chmin(x, y) x = min(x, y)
#define print(x) cout << x << "\\n"
#define print2(x, y) cout << x << " " << y << "\\n"
#define printv(vec)                                                            \\
  rep(lp, vec.size()) cout << vec[lp] << " ";                                  \\
  print("");
#define ALL(v) v.begin(), v.end()
#define SUM(v) accumulate(ALL(v), 0)
#define MAX(v) *max_element(ALL(v))
#define MIN(v) *min_element(ALL(v))
#define SORT(v) sort(ALL(v))
#define REV(v) reverse(ALL(v))
typedef long long ll;
using namespace std;
using vl = vector<ll>;
using vvl = vector<vector<ll>>;
using vs = vector<string>;

ll N, M, K;

const ll MAX = 510000;
long long fac[MAX], finv[MAX], inv[MAX];

void init() {
  fac[0] = fac[1] = 1;
  finv[0] = finv[1] = 1;
  inv[1] = 1;
  for (ll 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;
  }
}

ll cmb(ll n, ll k) {
  if (n < k)
    return 0;
  if (n < 0 || k < 0)
    return 0;
  return fac[n] * (finv[k] * finv[n - k] % MOD) % MOD;
}

ll modpow(ll a, ll n) {
  ll res = 1;
  while (n > 0) {
    if (n & 1)
      res = res * a % MOD;
    a = a * a % MOD;
    n >>= 1;
  }
  return res;
}

void Main() {

  cin >> N >> M >> K;
  ll ans = 0;
  init();
  rep(i, K + 1) {
    ll acc = M * cmb(N - 1, i) % MOD;
    acc *= modpow(M - 1, N - 1 - i);
    ans += acc;
    ans %= MOD;
  }
  print(ans);
}

int main() {

  ios_base::sync_with_stdio(false);
  cin.tie(0);
  Main();
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
#define INF 1000000007
#define rep(i, N) for (ll i = 0; i < N; i++)
#define rep2(i, j, k) for (ll i = j; i < k; i++)
#define chmax(x, y) x = max(x, y)
#define chmin(x, y) x = min(x, y)
#define print(x) cout << x << "\\n"
#define print2(x, y) cout << x << " " << y << "\\n"
#define printv(vec)                                                            \\
  rep(lp, vec.size()) cout << vec[lp] << " ";                                  \\
  print("");
#define ALL(v) v.begin(), v.end()
#define SUM(v) accumulate(ALL(v), 0)
#define MAX(v) *max_element(ALL(v))
#define MIN(v) *min_element(ALL(v))
#define SORT(v) sort(ALL(v))
#define REV(v) reverse(ALL(v))
typedef long long ll;
using namespace std;
using vl = vector<ll>;
using vvl = vector<vector<ll>>;
using vs = vector<string>;

ll MOD = 998244353;

ll modpow(ll a, ll n) {
  ll res = 1;
  while (n > 0) {
    if (n & 1)
      res = res * a % MOD;
    a = a * a % MOD;
    n >>= 1;
  }
  return res;
}

ll cmb(ll a, ll b) {
  if (b > a - b)
    return cmb(a, a - b);
  ll mul = 1;
  ll div = 1;
  for (ll i = 0; i < b; i++) {
    mul *= (a - i);
    mul %= MOD;
    div *= (i + 1);
    div %= MOD;
  }
  return mul * modpow(div, MOD - 2) % MOD;
}

void Main() {

  ll N, M, K;
  cin >> N >> M >> K;
  ll ans = 0;
  rep(i, K + 1) {
    ans += (M * cmb(N - 1, i) % MOD) * modpow(M - 1, N - 1 - i);
    ans %= MOD;
  }
  print(ans);
}

int main() {

  ios_base::sync_with_stdio(false);
  cin.tie(0);
  Main();
  return 0;
}