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

#include <bits/stdc++.h>

using namespace std;
using ll = long long;
using Graph = vector<vector<int>>;
using WGraph = vector<vector<pair<int, ll>>>;

template <class T> inline bool chmax(T &a, const T &b) {
  if (b > a) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}

constexpr int dx[4] = {-1, 0, 0, 1}, dy[4] = {0, -1, 1, 0};
constexpr int INF = 1e9;
constexpr int MOD = 1e9 + 7;
constexpr long long LINF = 1e18;
constexpr double EPS = 1e-10;
constexpr double PI = M_PI;

constexpr int mod = 998244353;

long long inv(long long n) {
  long long res = 1;
  while (n > 1) {
    res *= (mod - mod / n);
    res %= mod;
    n = mod % n;
  }
  return res;
}

void solve() {
  int n, m, k;
  cin >> n >> m >> k;

  vector<ll> powM(n);
  powM.at(0) = 1;
  for (int i = 1; i < n; ++i) {
    powM.at(i) = powM.at(i - 1) * (ll)(m - 1);
    powM.at(i) %= mod;
  }

  vector<ll> comb(n);
  ll nume = 1, deno = 1;
  for (int i = 0; i < n; ++i) {
    comb.at(i) = nume * inv(deno);
    comb.at(i) %= mod;
    nume *= (ll)(n - i - 1);
    nume %= mod;
    deno *= (ll)(i + 1);
    deno %= mod;
  }

  ll res = 0;
  for (ll i = 0; i <= k; ++i) {
    ll tmp = m;
    tmp *= comb.at(i);
    ;
    tmp %= mod;
    tmp *= powM.at(n - i - 1);
    tmp %= mod;
    res += tmp;
    res %= mod;
  }

  cout << res << '\\n';
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cout << fixed << setprecision(20);

  solve();

  return 0;
}


Error: TLE

Fixed submission:

#include <bits/stdc++.h>

using namespace std;
using ll = long long;
using Graph = vector<vector<int>>;
using WGraph = vector<vector<pair<int, ll>>>;

template <class T> inline bool chmax(T &a, const T &b) {
  if (b > a) {
    a = b;
    return true;
  }
  return false;
}
template <class T> inline bool chmin(T &a, const T &b) {
  if (b < a) {
    a = b;
    return true;
  }
  return false;
}

constexpr int dx[4] = {-1, 0, 0, 1}, dy[4] = {0, -1, 1, 0};
constexpr int INF = 1e9;
constexpr int MOD = 1e9 + 7;
constexpr long long LINF = 1e18;
constexpr double EPS = 1e-10;
constexpr double PI = M_PI;

constexpr int mod = 998244353;

long long inv(long long n) {
  long long res = 1;
  while (n > 1) {
    res *= (mod - mod / n);
    res %= mod;
    n = mod % n;
  }
  return res;
}

long long comb_mod(long long n, long long r) {
  long long nume = 1;
  long long deno = 1;
  long long r_new = (r <= n / 2) ? r : n - r;
  if (r_new < 0) {
    return 0;
  } else if (r_new == 0) {
    return 1;
  }

  for (long long i = 1; i <= r_new; i++) {
    deno *= i % mod;
    deno %= mod;
  }
  for (long long i = n - r_new + 1; i <= n; i++) {
    nume *= i % mod;
    nume %= mod;
  }
  deno = inv(deno);
  return (deno * nume) % mod;
}

void solve() {
  int n, m, k;
  cin >> n >> m >> k;

  vector<ll> powM(n);
  powM.at(0) = 1;
  for (int i = 1; i < n; ++i) {
    powM.at(i) = powM.at(i - 1) * (ll)(m - 1);
    powM.at(i) %= mod;
  }

  ll res = 0;
  for (ll i = 0; i <= k; ++i) {
    ll tmp = m;
    tmp *= comb_mod(n - 1, i);
    tmp %= mod;
    tmp *= powM.at(n - i - 1);
    tmp %= mod;
    res += tmp;
    res %= mod;
  }

  cout << res << '\\n';
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(0);
  cout << fixed << setprecision(20);

  solve();

  return 0;
}