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

#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
typedef unsigned long long int ull;
typedef pair<ll, ll> PLL;
#define chmax(x, y)                                                            \\
  do {                                                                         \\
    x = max(x, y);                                                             \\
  } while (0)
#define chmin(x, y)                                                            \\
  do {                                                                         \\
    x = min(x, y);                                                             \\
  } while (0)
#define _overload3(_1, _2, _3, name, ...) name
#define _rep(i, n) repi(i, 0, n)
#define repi(i, a, b)                                                          \\
  for (ll i = static_cast<ll>(a); i < static_cast<ll>(b); ++i)
#define rep(...) _overload3(__VA_ARGS__, repi, _rep, )(__VA_ARGS__) // NOLINT
template <typename T>
using reversed_priority_queue =
    std::priority_queue<T, std::vector<T>, std::greater<T>>;

// {{{ Modint
// using mint = modint<1000000007>;
// mint x = 2;
// std::cout << x.a << std::endl;
template <std::int64_t M> class modint {
public:
  int64_t a;

  explicit constexpr modint(const int64_t x = 0) noexcept
      : a((x % M + M) % M) {}
  constexpr int64_t &value() noexcept { return a; }
  constexpr const int64_t &value() const noexcept { return a; }
  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 {
    a += rhs.a;
    if (a >= M) {
      a -= M;
    }
    return *this;
  }
  constexpr modint &operator-=(const modint rhs) noexcept {
    if (a < rhs.a) {
      a += M;
    }
    a -= rhs.a;
    return *this;
  }
  constexpr modint &operator*=(const modint rhs) noexcept {
    a = a * rhs.a % M;
    return *this;
  }
  constexpr modint &operator/=(modint rhs) noexcept {
    int64_t exp = M - 2;
    while (exp) {
      if (exp % 2) {
        *this *= rhs;
      }
      rhs *= rhs;
      exp /= 2;
    }
    return *this;
  }
  constexpr modint pow(ll t) const {
    if (!t)
      return modint(1);
    modint a = pow(t >> 1);
    a *= a;
    if (t & 1)
      a *= *this;
    return a;
  }
};
// }}}
// 組み合わせO(n)のc4_init実行後、O(1)でc4を実行
std::vector<ll> fac(3000000, 1), finv(3000000, 1), inv(3000000, 1);

void c4_init(ll m) {
  for (ll i = 2; i < 3000000; i++) {
    fac[i] = fac[i - 1] * i % m;
    inv[i] = m - inv[m % i] * (m / i) % m;
    finv[i] = finv[i - 1] * inv[i] % m;
  }
}

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

using mint = modint<998244353>;
ll N, M, K;

signed main() {
  cin >> N >> M >> K;
  c4_init(998244353);
  mint ans;
  for (ll k = 0; k <= K; k++) {
    ans +=
        mint(M) * mint(c4(N - 1, k, 998244353)) * (mint(M - 1).pow(N - 1 - k));
  }
  cout << ans.value() << endl;
  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
typedef unsigned long long int ull;
typedef pair<ll, ll> PLL;
#define chmax(x, y)                                                            \\
  do {                                                                         \\
    x = max(x, y);                                                             \\
  } while (0)
#define chmin(x, y)                                                            \\
  do {                                                                         \\
    x = min(x, y);                                                             \\
  } while (0)
#define _overload3(_1, _2, _3, name, ...) name
#define _rep(i, n) repi(i, 0, n)
#define repi(i, a, b)                                                          \\
  for (ll i = static_cast<ll>(a); i < static_cast<ll>(b); ++i)
#define rep(...) _overload3(__VA_ARGS__, repi, _rep, )(__VA_ARGS__) // NOLINT
template <typename T>
using reversed_priority_queue =
    std::priority_queue<T, std::vector<T>, std::greater<T>>;

// {{{ Modint
// using mint = modint<1000000007>;
// mint x = 2;
// std::cout << x.a << std::endl;
template <std::int64_t M> class modint {
public:
  int64_t a;

  explicit constexpr modint(const int64_t x = 0) noexcept
      : a((x % M + M) % M) {}
  constexpr int64_t &value() noexcept { return a; }
  constexpr const int64_t &value() const noexcept { return a; }
  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 {
    a += rhs.a;
    if (a >= M) {
      a -= M;
    }
    return *this;
  }
  constexpr modint &operator-=(const modint rhs) noexcept {
    if (a < rhs.a) {
      a += M;
    }
    a -= rhs.a;
    return *this;
  }
  constexpr modint &operator*=(const modint rhs) noexcept {
    a = a * rhs.a % M;
    return *this;
  }
  constexpr modint &operator/=(modint rhs) noexcept {
    int64_t exp = M - 2;
    while (exp) {
      if (exp % 2) {
        *this *= rhs;
      }
      rhs *= rhs;
      exp /= 2;
    }
    return *this;
  }
  constexpr modint pow(ll t) const {
    if (!t)
      return modint(1);
    modint a = pow(t >> 1);
    a *= a;
    if (t & 1)
      a *= *this;
    return a;
  }
};
// }}}

using mint = modint<998244353>;
ll N, M, K;
map<PLL, mint> memo;

mint f(ll n, ll k) {
  auto it = memo.find(PLL(n, k));
  if (it != memo.end())
    return it->second;

  mint res;
  if (n == 1) {
    return mint(1);
  }

  res = mint(M - 1) * f(n - 1, k);
  if (k > 0)
    res += f(n - 1, k - 1);

  return res;
}

signed main() {
  cin >> N >> M >> K;
  auto ans = mint(M) * f(N, K);
  cout << ans.value() << endl;
  return 0;
}