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

#include <bits/stdc++.h>

using namespace std;
using ll = long long;
using vl = vector<ll>;
using vvl = vector<vl>;
using pll = pair<ll, ll>;

const ll oo = 0x3f3f3f3f3f3f3f3fLL;
const double eps = 1e-9;

#define sz(c) ll((c).size())
#define all(c) begin(c), end(c)
#define mp make_pair
#define mt make_tuple
#define pb push_back
#define eb emplace_back
#define xx first
#define yy second
#define FOR(i, a, b) for (ll i = (a); i < (b); i++)
#define FORD(i, a, b) for (ll i = ll(b) - 1; i >= (a); i--)
#define TR(X)                                                                  \\
  ({                                                                           \\
    if (1)                                                                     \\
      cerr << "TR: " << (#X) << " = " << (X) << endl;                          \\
  })

constexpr ll MOD = 998244353ll;

template <long long mod = 1000000007> struct modint {
  long long a;

  modint() : a(0) {}
  modint(long long t) {
    a = t % mod;
    if (a < 0)
      a += mod;
  }

  operator long long() const { return a; }

  bool operator==(const modint &x) const { return a == x.a; }
  bool operator!=(const modint &x) const { return a != x.a; }

  modint operator-() const { return modint(a ? (mod - a) : 0); }
  modint operator~() const { return pow(mod - 2); }

  modint operator+(const modint &x) const { return modint(*this) += x; }
  modint operator-(const modint &x) const { return modint(*this) -= x; }
  modint operator*(const modint &x) const { return modint(*this) *= x; }
  modint operator/(const modint &x) const { return modint(*this) /= x; }

  modint &operator+=(const modint &x) {
    a += x.a;
    if (a >= mod)
      a -= mod;
    return *this;
  }
  modint &operator-=(const modint &x) {
    a -= x.a;
    if (a < 0)
      a += mod;
    return *this;
  }
  modint &operator*=(const modint &x) {
    a = a * x.a % mod;
    return *this;
  }
  modint &operator/=(const modint &x) {
    a = a * (~x).a % mod;
    return *this;
  }

  friend ostream &operator<<(ostream &os, const modint &x) { return os << x.a; }
  friend istream &operator>>(istream &is, modint &x) {
    long long t;
    is >> t;
    x = modint(t);
    return is;
  }

  modint pow(long long x) const {
    modint ret = 1, tmp = a;
    for (; x; tmp *= tmp, x >>= 1) {
      if (x & 1)
        ret *= tmp;
    }
    return ret;
  }
};

// modint<MOD> binom(ll n, ll k) {
//     modint<MOD> res = 1;
//     FOR(i, 0, k) {
//         res *= n - i;
//         res /= i+1;
//     }
//     return res;
// }

int main() {
  cin.sync_with_stdio(0);
  ll N, M, K;
  cin >> N >> M >> K;

  vector<modint<MOD>> fac(N + 5), facInv(N + 5);
  fac[0] = 1;
  FOR(i, 1, N + 5) fac[i] = fac[i - 1] * modint<MOD>(i);
  facInv[N + 4] = ~fac[N + 4];
  FORD(i, 1, N + 5) facInv[i - 1] = facInv[i] * modint<MOD>(i);

  auto binom = [&](ll n, ll k) {
    if (k < 0 || n < k)
      return modint<MOD>(0);
    return fac[n] * facInv[k] * facInv[n - k];
  };

  modint<MOD> ans = 0;

  FOR(i, 0, K + 1) {
    ans += binom(N - 1, i) * modint<MOD>(M) * modint<MOD>(M - 1).pow(N - i - 1);
  }
  cout << ans << endl;
} // cin.tie(0) bei schnellem Wechseln


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;
using ll = long long;
using vl = vector<ll>;
using vvl = vector<vl>;
using pll = pair<ll, ll>;

const ll oo = 0x3f3f3f3f3f3f3f3fLL;
const double eps = 1e-9;

#define sz(c) ll((c).size())
#define all(c) begin(c), end(c)
#define mp make_pair
#define mt make_tuple
#define pb push_back
#define eb emplace_back
#define xx first
#define yy second
#define FOR(i, a, b) for (ll i = (a); i < (b); i++)
#define FORD(i, a, b) for (ll i = ll(b) - 1; i >= (a); i--)
#define TR(X)                                                                  \\
  ({                                                                           \\
    if (1)                                                                     \\
      cerr << "TR: " << (#X) << " = " << (X) << endl;                          \\
  })

constexpr ll MOD = 998244353ll;

template <long long mod = 1000000007> struct modint {
  long long a;

  modint() : a(0) {}
  modint(long long t) {
    a = t % mod;
    if (a < 0)
      a += mod;
  }

  operator long long() const { return a; }

  bool operator==(const modint &x) const { return a == x.a; }
  bool operator!=(const modint &x) const { return a != x.a; }

  modint operator-() const { return modint(a ? (mod - a) : 0); }
  modint operator~() const { return pow(mod - 2); }

  modint operator+(const modint &x) const { return modint(*this) += x; }
  modint operator-(const modint &x) const { return modint(*this) -= x; }
  modint operator*(const modint &x) const { return modint(*this) *= x; }
  modint operator/(const modint &x) const { return modint(*this) /= x; }

  modint &operator+=(const modint &x) {
    a += x.a;
    if (a >= mod)
      a -= mod;
    return *this;
  }
  modint &operator-=(const modint &x) {
    a -= x.a;
    if (a < 0)
      a += mod;
    return *this;
  }
  modint &operator*=(const modint &x) {
    a = a * x.a % mod;
    return *this;
  }
  modint &operator/=(const modint &x) {
    a = a * (~x).a % mod;
    return *this;
  }

  friend ostream &operator<<(ostream &os, const modint &x) { return os << x.a; }
  friend istream &operator>>(istream &is, modint &x) {
    long long t;
    is >> t;
    x = modint(t);
    return is;
  }

  modint pow(long long x) const {
    modint ret = 1, tmp = a;
    for (; x; tmp *= tmp, x >>= 1) {
      if (x & 1)
        ret *= tmp;
    }
    return ret;
  }
};

modint<MOD> binom(ll n, ll k) {
  modint<MOD> res = 1;
  FOR(i, 0, k) {
    res *= n - i;
    res /= i + 1;
  }
  return res;
}

int main() {
  cin.sync_with_stdio(0);
  ll N, M, K;
  cin >> N >> M >> K;

  modint<MOD> ans = 0;

  FOR(i, 0, K + 1) {
    ans += binom(N - 1, i) * modint<MOD>(M) * modint<MOD>(M - 1).pow(N - i - 1);
  }
  cout << ans << endl;
} // cin.tie(0) bei schnellem Wechseln