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

#include <algorithm>
#include <bitset>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <string>
#include <tuple>
#include <utility>
#include <vector>

using ll = long long;
using ld = long double;
using namespace std;
const int INF = 1e9 + 100;
const ll INF64 = 7e18l;
const int mod = 998244353;
#define rep(i, a, b) for (int i = (a); i < (b); i++)
#define ALL(a) (a).begin(), (a).end()
#define FIN                                                                    \\
  ios::sync_with_stdio(0);                                                     \\
  cin.tie(0);                                                                  \\
  cout.tie(0);

struct mint {
  ll val;
  constexpr mint(ll v = 0) noexcept : val(v % mod) {
    if (val < 0)
      val += mod;
  }
  constexpr mint operator+(const mint &rhs) const noexcept {
    return mint(*this) += rhs;
  }
  constexpr mint operator-(const mint &rhs) const noexcept {
    return mint(*this) -= rhs;
  }
  constexpr mint operator*(const mint &rhs) const noexcept {
    return mint(*this) *= rhs;
  }
  constexpr mint operator/(const mint &rhs) const noexcept {
    return mint(*this) /= rhs;
  }
  constexpr mint operator~() const noexcept {
    ll a = this->val, b = mod, u = 1, v = 0;
    while (b > 0) {
      ll t = a / b;
      a -= t * b;
      u -= t * v;
      ll c = a;
      a = b;
      b = c;
      ll w = u;
      u = v;
      v = w;
    }
    if (u < 0)
      u += mod;
    u %= mod;
    return u;
  }
  constexpr mint &operator+=(const mint &rhs) noexcept {
    val += rhs.val;
    if (val >= mod)
      val -= mod;
    return *this;
  }
  constexpr mint &operator-=(const mint &rhs) noexcept {
    if (val < rhs.val)
      val += mod;
    val -= rhs.val;
    return *this;
  }
  constexpr mint &operator*=(const mint &rhs) noexcept {
    val = val * rhs.val % mod;
    return *this;
  }
  constexpr mint &operator/=(const mint &rhs) noexcept {
    val = val * (~(mint(rhs))).val % mod;
    return *this;
  }
  friend mint modpow(mint a, ll n) noexcept {
    auto res = mint(1);
    while (n > 0) {
      if (n & 1)
        res *= a;
      a *= a;
      n >>= 1;
    }
    return res;
  }
  friend std::ostream &operator<<(std::ostream &os, const mint &a) noexcept {
    return os << a.val;
  }
  friend std::string debug_s(mint a) noexcept { return std::to_string(a.val); }
};
namespace combinatorics {
const int limit = 2000020;
bool initialized = false;
long long fac_table[limit], inv_table[limit], finv_table[limit];
void initialize() {
  fac_table[0] = fac_table[1] = 1;
  inv_table[1] = 1;
  finv_table[0] = finv_table[1] = 1;
  for (int i = 2; i < limit; i++) {
    fac_table[i] = fac_table[i - 1] * i % mod;
    inv_table[i] = mod - inv_table[mod % i] * (mod / i) % mod;
    finv_table[i] = finv_table[i - 1] * inv_table[i] % mod;
  }
  initialized = true;
}
long long com(int n, int k) {
  if (!initialized)
    initialize();
  if (n < k || n < 0 || k < 0)
    return 0;
  else
    return fac_table[n] * (finv_table[k] * finv_table[n - k] % mod) % mod;
}
long long perm(int n, int k) {
  if (!initialized)
    initialize();
  if (n < k || n < 0 || k < 0)
    return 0;
  else
    return fac_table[n] * finv_table[n - k] % mod;
}
} // namespace combinatorics
using combinatorics::com;
using combinatorics::perm;

int main() {
  ll n, m, k;
  cin >> n >> m >> k;

  mint ans = 0;
  for (ll i = 0; i <= k; i++) {
    ans += mint(m) * com(n - 1, i) * modpow(mint(m - 1), n - 1 - i);
    // m じゃなくて、mint(m)にしないとできない。
  }

  cout << ans << endl;
  return 0;
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <bitset>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <list>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <string>
#include <tuple>
#include <utility>
#include <vector>

using ll = long long;
using ld = long double;
using namespace std;
const int INF = 1e9 + 100;
const ll INF64 = 7e18l;
const int mod = 998244353;
#define rep(i, a, b) for (int i = (a); i < (b); i++)
#define ALL(a) (a).begin(), (a).end()
#define FIN                                                                    \\
  ios::sync_with_stdio(0);                                                     \\
  cin.tie(0);                                                                  \\
  cout.tie(0);

template <int MOD> struct Fp {
  long long val;
  constexpr Fp(long long v = 0) noexcept : val(v % MOD) {
    if (val < 0)
      val += MOD;
  }
  constexpr int getmod() { return MOD; }
  constexpr Fp operator-() const noexcept { return val ? MOD - val : 0; }
  constexpr Fp operator+(const Fp &r) const noexcept { return Fp(*this) += r; }
  constexpr Fp operator-(const Fp &r) const noexcept { return Fp(*this) -= r; }
  constexpr Fp operator*(const Fp &r) const noexcept { return Fp(*this) *= r; }
  constexpr Fp operator/(const Fp &r) const noexcept { return Fp(*this) /= r; }
  constexpr Fp &operator+=(const Fp &r) noexcept {
    val += r.val;
    if (val >= MOD)
      val -= MOD;
    return *this;
  }
  constexpr Fp &operator-=(const Fp &r) noexcept {
    val -= r.val;
    if (val < 0)
      val += MOD;
    return *this;
  }
  constexpr Fp &operator*=(const Fp &r) noexcept {
    val = val * r.val % MOD;
    return *this;
  }
  constexpr Fp &operator/=(const Fp &r) noexcept {
    long long a = r.val, b = MOD, u = 1, v = 0;
    while (b) {
      long long t = a / b;
      a -= t * b;
      swap(a, b);
      u -= t * v;
      swap(u, v);
    }
    val = val * u % MOD;
    if (val < 0)
      val += MOD;
    return *this;
  }
  constexpr bool operator==(const Fp &r) const noexcept {
    return this->val == r.val;
  }
  constexpr bool operator!=(const Fp &r) const noexcept {
    return this->val != r.val;
  }
  friend constexpr ostream &operator<<(ostream &os, const Fp<MOD> &x) noexcept {
    return os << x.val;
  }
  // a^n
  friend constexpr Fp<MOD> modpow(const Fp<MOD> &a, long long n) noexcept {
    if (n == 0)
      return 1;
    auto t = modpow(a, n / 2);
    t = t * t;
    if (n & 1)
      t = t * a;
    return t;
  }
};
using mint = Fp<998244353>;

mint combination(ll n, ll k) {
  mint res = 1;
  for (ll i = n; i > n - k; i--)
    res = res * i;
  for (ll i = k; i > 0; i--)
    res = res / i;
  return res;
}

int main() {
  FIN;
  ll n, m, k;
  cin >> n >> m >> k;

  mint ans = 0;
  for (ll i = 0; i <= k; i++) {
    mint cur = m;
    cur *= combination(n - 1, i);
    for (ll j = 0; j < n - i - 1; j++)
      cur *= (m - 1);
    ans += cur;
  }

  cout << ans << endl;
  return 0;
}