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

#include <bits/stdc++.h>
using namespace std;
#define mod 998244353
typedef long long ll;

long long fact(int N) {
  long long ans = 1;
  for (int i = 1; i <= N; i++)
    ans = ((ans % mod) * i) % mod;
  return ans % mod;
}

long long fast_exp(ll a, int b) {
  long long x = a;
  long long y = (b % 2) ? x : 1;
  b >>= 1;
  while (b > 0) {
    x = ((x % mod) * (x % mod)) % mod;
    if (b % 2 == 1) {
      y = (y == 1) ? x : (((y % mod) * (x % mod)) % mod);
    }
    b >>= 1;
  }
  return y;
}
struct mint {
  ll x;
  mint(ll x = 0) : x(x % mod) {}
  mint &operator+=(const mint a) {
    if ((x += a.x) >= mod)
      x -= mod;
    return *this;
  }
  mint &operator-=(const mint a) {
    if ((x += mod - a.x) >= mod)
      x -= mod;
    return *this;
  }
  mint &operator*=(const mint a) {
    (x *= a.x) %= mod;
    return *this;
  }
  mint operator+(const mint a) const {
    mint res(*this);
    return res += a;
  }
  mint operator-(const mint a) const {
    mint res(*this);
    return res -= a;
  }
  mint operator*(const mint a) const {
    mint res(*this);
    return res *= a;
  }
};
int M;
mint nCr(int n, int r) {
  r = min(n - r, r);
  mint up = 1;
  mint down = 1;
  for (int i = 0; i < r; i++) {
    up *= (n - i);
    down *= (i + 1);
  }
  mint ans = fast_exp(down.x, mod - 2);
  ans *= up;
  return ans;
}

int main() {
  int n, m, k;
  cin >> n >> m >> k;
  mint ans1 = 0;

  vector<mint> co(k + 1);
  co[k] = fast_exp(m - 1, n - 1 - k);
  for (int i = k - 1; i >= 0; i--) {
    co[i] = co[i + 1] * (m - 1);
  }
  if (m == 1 && k == n - 1) {
    cout << 1 << endl;
    return 0;
  }

  vector<mint> inv(n);
  inv[n - 1] = fast_exp(fact(n - 1), mod - 2);
  for (int i = n - 2; i >= 1; i--)
    inv[i] = inv[i + 1] * (i + 1);
  inv[0] = 1;
  mint u1 = fact(n - 1);
  for (int i = 0; i <= k; i++) {
    mint t = co[i] * u1 * inv[n - 1 - i] * inv[i];
    ans1 += t;
  }
  ans1 *= m;
  cout << ans1.x << endl;

  return 0;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
#define mod 998244353
typedef long long ll;

long long fact(int N) {
  long long ans = 1;
  for (int i = 1; i <= N; i++)
    ans = ((ans % mod) * i) % mod;
  return ans % mod;
}

long long fast_exp(ll a, int b) {
  long long x = a;
  long long y = (b % 2) ? x : 1;
  b >>= 1;
  while (b > 0) {
    x = ((x % mod) * (x % mod)) % mod;
    if (b % 2 == 1) {
      y = (y == 1) ? x : (((y % mod) * (x % mod)) % mod);
    }
    b >>= 1;
  }
  return y;
}
struct mint {
  ll x;
  mint(ll x = 0) : x(x % mod) {}
  mint &operator+=(const mint a) {
    if ((x += a.x) >= mod)
      x -= mod;
    return *this;
  }
  mint &operator-=(const mint a) {
    if ((x += mod - a.x) >= mod)
      x -= mod;
    return *this;
  }
  mint &operator*=(const mint a) {
    (x *= a.x) %= mod;
    return *this;
  }
  mint operator+(const mint a) const {
    mint res(*this);
    return res += a;
  }
  mint operator-(const mint a) const {
    mint res(*this);
    return res -= a;
  }
  mint operator*(const mint a) const {
    mint res(*this);
    return res *= a;
  }
};
int M;
mint nCr(int n, int r) {
  r = min(n - r, r);
  mint up = 1;
  mint down = 1;
  for (int i = 0; i < r; i++) {
    up *= (n - i);
    down *= (i + 1);
  }
  mint ans = fast_exp(down.x, mod - 2);
  ans *= up;
  return ans;
}

int main() {
  int n, m, k;
  cin >> n >> m >> k;
  mint ans = 0;
  vector<mint> co(k + 1);
  co[k] = fast_exp(m - 1, n - 1 - k);

  for (int i = k - 1; i >= 0; i--) {
    co[i] = co[i + 1] * (m - 1);
  }
  mint x = 1;
  for (int i = 0; i <= k; i++) {
    mint t = co[i] * nCr(n - 1, i);
    ans += t;
  }
  ans *= m;
  cout << ans.x << endl;
  return 0;
}