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

/*COMPETITIVE PROGRAMMING C++ TEMPLATE */

#include <algorithm>
#include <climits>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#define all(x) x.begin(), x.end()
#define MOD 998244353

const long double PI = 3.141592653589793236L;

typedef long long int ll;
typedef long double ld;

using namespace std;

map<pair<int, int>, ll> mp;
const int maxN = 2e5;
vector<ll> fact(maxN + 1, 1);

void calc_fact() {
  for (int i = 1; i <= maxN; i++) {
    fact[i] = fact[i - 1] * i;
    fact[i] %= MOD;
  }
}

ll C(ll n, ll r) {
  if (n == 0) {
    return 0;
  }
  if (n == r || r == 0) {
    return 1;
  }
  if (mp[{n, r}]) {
    return mp[{n, r}];
  }

  ll x, y;

  if (mp[{n - 1, r - 1}]) {
    x = mp[{n - 1, r - 1}];
  } else {
    x = C(n - 1, r - 1);
    if (x >= MOD)
      x %= MOD;
    mp[{n - 1, r - 1}] = x;
  }

  if (mp[{n - 1, r}]) {
    y = mp[{n - 1, r}];
  } else {
    y = C(n - 1, r);
    if (y >= MOD)
      y %= MOD;
    mp[{n - 1, r}] = y;
  }

  ll ret = mp[{n, r}] = x + y;

  if (ret >= MOD)
    ret %= MOD;

  return ret;
}

ll power(ll X, ll N) {
  ll res = 1;

  while (N > 0) {
    if (N & 1) {
      res *= X;
      if (res >= MOD)
        res %= MOD;
    }

    X = X * X;
    if (X >= MOD)
      X %= MOD;

    N = N >> 1;
  }

  return res;
}

ll c(ll n, ll r) {
  ll num = fact[n];
  ll den = fact[n - r] * fact[r];
  den %= MOD;

  den = power(den, MOD - 2);

  return (num * den) % MOD;
}

ll calc(ll N, ll M, ll K) {
  ll ans = 1;
  ans *= M;

  ans *= power(M - 1, N - K - 1);
  if (ans >= MOD)
    ans %= MOD;

  ans *= c(N - 1, K);

  if (ans >= MOD)
    ans %= MOD;

  return ans;
}

void solve() {
  ll N, M, K;
  cin >> N >> M >> K;

  ll ans = 0;

  for (int i = 0; i <= K; i++) {
    ans += calc(N, M, i);
    if (ans >= MOD)
      ans %= MOD;
  }
  cout << ans << endl;
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);

  calc_fact();
  int T = 1;
  // cin >> T;

  while (T--) {
    solve();
  }

  return 0;
}


Error: TLE

Fixed submission:

/*COMPETITIVE PROGRAMMING C++ TEMPLATE */

#include <algorithm>
#include <climits>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>

#define all(x) x.begin(), x.end()
#define MOD 998244353

const long double PI = 3.141592653589793236L;

typedef long long int ll;
typedef long double ld;

using namespace std;

map<pair<int, int>, ll> mp;

ll C(ll n, ll r) {
  if (n == r || r == 0) {
    return 1;
  }
  if (mp[{n, r}]) {
    return mp[{n, r}];
  }

  return mp[{n, r}] = (C(n - 1, r - 1) % MOD + C(n - 1, r) % MOD) % MOD;
}

ll power(ll X, ll N) {
  ll res = 1;

  while (N > 0) {
    if (N & 1) {
      res *= X;
      res %= MOD;
    }

    X = X * X;
    X %= MOD;

    N = N >> 1;
  }

  return res;
}

ll calc(ll N, ll M, ll K) {
  ll ans = 1;
  ans *= M;

  ans *= power(M - 1, N - K - 1);
  ans %= MOD;

  ans *= C(N - 1, K);
  ans %= MOD;

  return ans;
}

void solve() {
  ll N, M, K;
  cin >> N >> M >> K;

  ll ans = 0;

  for (int i = 0; i <= K; i++) {
    ans += calc(N, M, i);
    ans %= MOD;
  }
  cout << ans << endl;
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);

  int T = 1;
  // cin >> T;

  while (T--) {
    solve();
  }

  return 0;
}