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

#include <cstdint>
#include <iostream>
#include <vector>

using namespace std;
using ll = long long;

#define ALL(x) begin(x), end(x)
#define REP(i, n) for (size_t i = 0, i##_len = (n); i < i##_len; ++i)

template <std::uint_fast64_t Mod> class ModInt {
  using u64 = std::uint_fast64_t;

public:
  u64 val;

  constexpr ModInt(const u64 x = 0) noexcept : val(x % Mod) {}
  constexpr ModInt operator+(const ModInt rhs) const noexcept {
    return ModInt(*this) += rhs;
  }
  constexpr ModInt operator-(const ModInt rhs) const noexcept {
    return ModInt(*this) -= rhs;
  }
  constexpr ModInt operator*(const ModInt rhs) const noexcept {
    return ModInt(*this) *= rhs;
  }
  constexpr ModInt operator/(const ModInt rhs) const noexcept {
    return ModInt(*this) /= rhs;
  }
  constexpr ModInt &operator+=(const ModInt rhs) noexcept {
    val += rhs.val;
    if (val >= Mod) {
      val -= Mod;
    }
    return *this;
  }
  constexpr ModInt &operator-=(const ModInt rhs) noexcept {
    if (val < rhs.val) {
      val += Mod;
    }
    val -= rhs.val;
    return *this;
  }
  constexpr ModInt &operator*=(const ModInt rhs) noexcept {
    val = val * rhs.val % Mod;
    return *this;
  }
  constexpr ModInt &operator/=(ModInt rhs) noexcept {
    u64 exp = Mod - 2;
    while (exp) {
      if (exp % 2) {
        *this *= rhs;
      }
      rhs *= rhs;
      exp /= 2;
    }
    return *this;
  }
};

template <typename T> T power(T x, int n) {
  T ans{1};
  while (n) {
    if (n % 2)
      ans *= x;
    x *= x;
    n >>= 1;
  }
  return ans;
}

const ll MOD = 998244353;

using mint = ModInt<MOD>;

void solve(ll N, ll M, ll K) {
  mint ans = 0;
  mint nck = 1;
  for (auto i = 0; i <= K; ++i) {
    ans += nck * power<mint>(M - 1, N - i - 1);
    nck *= N - 1 - i;
    nck /= i + 1;
  }
  ans *= M;

  cout << ans.val << endl;
}

int main() {
  ll N;
  cin >> N;
  ll M;
  cin >> M;
  ll K;
  cin >> K;
  solve(N, M, K);
  return 0;
}


Error: TLE

Fixed submission:

#include <cstdint>
#include <iostream>
#include <vector>

using namespace std;
using ll = long long;

#define ALL(x) begin(x), end(x)
#define REP(i, n) for (size_t i = 0, i##_len = (n); i < i##_len; ++i)

template <std::uint_fast64_t Mod> class ModInt {
  using u64 = std::uint_fast64_t;

public:
  u64 val;

  constexpr ModInt(const u64 x = 0) noexcept : val(x % Mod) {}
  constexpr ModInt operator+(const ModInt rhs) const noexcept {
    return ModInt(*this) += rhs;
  }
  constexpr ModInt operator-(const ModInt rhs) const noexcept {
    return ModInt(*this) -= rhs;
  }
  constexpr ModInt operator*(const ModInt rhs) const noexcept {
    return ModInt(*this) *= rhs;
  }
  constexpr ModInt operator/(const ModInt rhs) const noexcept {
    return ModInt(*this) /= rhs;
  }
  constexpr ModInt &operator+=(const ModInt rhs) noexcept {
    val += rhs.val;
    if (val >= Mod) {
      val -= Mod;
    }
    return *this;
  }
  constexpr ModInt &operator-=(const ModInt rhs) noexcept {
    if (val < rhs.val) {
      val += Mod;
    }
    val -= rhs.val;
    return *this;
  }
  constexpr ModInt &operator*=(const ModInt rhs) noexcept {
    val = val * rhs.val % Mod;
    return *this;
  }
  constexpr ModInt &operator/=(ModInt rhs) noexcept {
    u64 exp = Mod - 2;
    while (exp) {
      if (exp % 2) {
        *this *= rhs;
      }
      rhs *= rhs;
      exp /= 2;
    }
    return *this;
  }
};

#include <algorithm>

template <typename T> T nCk(int n, int k) {
  if (k > n)
    return 0;
  T ret = 1;
  for (auto i = 0; i < std::min(n - k, k); ++i) {
    ret *= n - i;
    ret /= i + 1;
  }
  return ret;
}

template <typename T> T power(T x, int n) {
  T ans{1};
  while (n) {
    if (n % 2)
      ans *= x;
    x *= x;
    n >>= 1;
  }
  return ans;
}

const ll MOD = 998244353;

using mint = ModInt<MOD>;

void solve(ll N, ll M, ll K) {
  mint ans = 0;
  for (auto i = 0; i <= K; ++i) {
    ans += nCk<mint>(N - 1, i) * power<mint>(M - 1, N - i - 1);
  }
  ans *= M;

  cout << ans.val << endl;
}

int main() {
  ll N;
  cin >> N;
  ll M;
  cin >> M;
  ll K;
  cin >> K;
  solve(N, M, K);
  return 0;
}