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

#include <bits/stdc++.h>
using namespace std;
static const long long MOD = 1000000000 + 7;
static const double PI = 3.141592653589793;
vector<long long> vecmodfact(2, 1);

long long modpow(long long x, long long n, long long mod = 0) {
  if (mod == 0) {
    long long res = 1;
    while (n > 0) {
      if (n % 2 == 1) {
        res = (res * x);
      }
      x = (x * x);
      n /= 2;
    }
    return res;
  }

  long long res = 1;
  while (n > 0) {
    if (n % 2 == 1) {
      res = (res * x) % mod;
    }
    x = (x * x) % mod;
    n /= 2;
  }
  return res;
}

long long modinverse(long long x, long long mod) {
  // mod must be prime number.
  return modpow(x, mod - 2, mod);
}

long long modfact(long long n, long long mod) {
  if (vecmodfact.size() - 1 >= n)
    return vecmodfact.at(n);
  int i = vecmodfact.size();
  vecmodfact.resize(n + 1);
  for (; i <= vecmodfact.size() - 1; i++) {
    vecmodfact.at(i) = vecmodfact.at(i - 1) * i % mod;
  }
  return vecmodfact.at(n);
}

long long modcomb(long long m, long long n, long long mod) {
  if (m < n || n < 0)
    return 0;
  if (n == 0 || n == m)
    return 1;
  long long ans = modfact(m, mod) * modinverse(modfact(n, mod), mod) % mod;
  ans = ans * modinverse(modfact(m - n, mod), mod) % mod;
  return ans;
}

long long solve(long long n, long long m, long long k, long long mod) {
  long long ans = m * modcomb(n - 1, k, mod) % mod;
  ans = ans * modpow(m - 1, n - k - 1, mod) % mod;
  return ans;
}

int main() {
  long long mod = 998244353;
  long long n, m, k;
  cin >> n >> m >> k;

  long long ans = 0;
  for (int K = 0; K <= k; K++) {
    ans = (ans + solve(n, m, K, mod)) % mod;
  }

  cout << ans << endl;
}

Error: TLE

Fixed submission:

#include <bits/stdc++.h>
using namespace std;
static const long long MOD = 1000000000 + 7;
static const double PI = 3.141592653589793;

long long modpow(long long x, long long n, long long mod = 0) {
  if (mod == 0) {
    long long res = 1;
    while (n > 0) {
      if (n % 2 == 1) {
        res = (res * x);
      }
      x = (x * x);
      n /= 2;
    }
    return res;
  }

  long long res = 1;
  while (n > 0) {
    if (n % 2 == 1) {
      res = (res * x) % mod;
    }
    x = (x * x) % mod;
    n /= 2;
  }
  return res;
}

long long modinverse(long long x, long long mod) {
  // mod must be prime number.
  return modpow(x, mod - 2, mod);
}

long long modcomb(long long m, long long n, long long mod) {
  if (m < n)
    return 0;
  if (n > m / 2)
    n = m - n;
  if (n == 0)
    return 1;
  long long m1, m2, m3;
  long long fact = 1;
  for (int i = 1; i <= m; i++) {
    fact = fact * i % mod;
    if (i == m)
      m1 = fact;
    if (i == n)
      m2 = fact;
    if (i == m - n)
      m3 = fact;
  }
  long long ans = m1 * modinverse(m2, mod) % mod;
  ans = ans * modinverse(m3, mod) % mod;
  return ans;
}

long long solve(long long n, long long m, long long k, long long mod) {
  long long ans = m * modcomb(n - 1, k, mod) % mod;
  ans = ans * modpow(m - 1, n - k - 1, mod) % mod;
  return ans;
}

int main() {
  long long mod = 998244353;
  long long n, m, k;
  cin >> n >> m >> k;

  long long ans = 0;
  for (int K = 0; K <= k; K++) {
    ans = (ans + solve(n, m, K, mod)) % mod;
  }

  cout << ans << endl;
}