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

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#define ll long long
#define fr(i, k, N) for (decltype(N + k) i = k; i < N; i++)
#define rep(i, N) for (decltype(N) i = 0; i < N; i++)
#define ALL(v) v.begin(), v.end()
using namespace std;
template <class X> void pr(X test) { cout << test << endl; }
template <class X> void prr(X test) {
  for (auto it : test)
    cout << it << " ";
  cout << endl;
}
template <class X> void prrr(X test) {
  int f = 0;
  for (auto it : test)
    cout << (f++ != 0 ? " " : "") << it;
  cout << endl;
}
template <typename X> istream &operator>>(istream &i, vector<X> &v) {
  for (X &x : v)
    i >> x;
  return i;
}

class UF { // union find from https://github.com/kartikkukreja/blog-codes (MIT)
           // [usage]
           // auto uf = new UF(n);
           // uf->merge(l, r); // 0 <= l, r <= n-1
           // uf->sz[uf->find(l)]; // size of l's union
           // ...
public:
  int *id; // inner id
  int cnt; // number of all vartices
  int *sz; // sz[i] means size of i's set
  UF(int N) {
    cnt = N;
    id = new int[N];
    sz = new int[N];
    for (int i = 0; i < N; i++) {
      id[i] = i;
      sz[i] = 1;
    }
  }
  ~UF() {
    delete[] id;
    delete[] sz;
  }
  int find(int p) {
    int root = p;
    while (root != id[root])
      root = id[root];
    while (p != root) {
      int newp = id[p];
      id[p] = root;
      p = newp;
    }
    return root;
  }
  void merge(int x, int y) {
    int i = find(x);
    int j = find(y);
    if (i == j)
      return;

    if (sz[i] < sz[j]) {
      id[i] = j;
      sz[j] += sz[i];
    } else {
      id[j] = i;
      sz[i] += sz[j];
    }
    cnt--;
  }
  bool connected(int x, int y) { return find(x) == find(y); }
  int count() { return cnt; }
};

// return 約数unordered_set 1 <= x <= n
auto divv(ll n) {
  unordered_set<decltype(n)> v;
  for (decltype(n) i = 1; i * i <= n; i++) {
    if (n % i == 0) {
      v.insert(i);
      if (i * i != n)
        v.insert(n / i);
    }
  }
  return v;
};

// check prime (and n>=2) or not
auto isPrime(ll n) {
  if (n <= 1)
    return false;
  for (decltype(n) i = 2; i <= sqrt(n); ++i) {
    if (n % i == 0) {
      return false;
    }
  }
  return true;
}

// 最小公倍数
auto lcm_(ll a, ll b) {
  if (a < b)
    swap(a, b);
  auto r = a % b;
  auto x = a * b;
  while (r != 0) {
    a = b;
    b = r;
    r = a % b;
  }
  return x / b;
}

// 最大公約数
auto gcd_(ll a, ll b) { return (a * b) / lcm_(a, b); };

// factorial n mod prm
ll factmod(ll a, ll prm) {
  return a == 0 ? 1 : (a * factmod(a - 1, prm)) % prm;
};

// 1/x mod prm
auto invmod = [](auto x, auto prm) {
  decltype(x) r = 1;
  if (x == 1)
    return r;
  for (auto ex = prm - 2; ex; ex >>= 1) {
    if (ex & 1)
      r = (r * x) % prm;
    x = (x * x) % prm;
  }
  return r;
};

// {{2, 3}, {3, 1}, {5, 2}, ...}
unordered_map<ll, ll> prime_factor(ll n) {
  unordered_map<ll, ll> ret;
  for (ll i = 2; i * i <= n; i++) {
    while (n % i == 0) {
      ret[i]++;
      n /= i;
    }
  }
  if (n != 1)
    ret[n] = 1;
  return ret;
}

ll modpow(ll base, ll exp, ll mod) {
  base %= mod;
  ll result = 1;
  while (exp > 0) {
    if (exp & 1)
      result = (result * base) % mod;
    base = (base * base) % mod;
    exp >>= 1;
  }
  return result;
}

// get {1, 1, 2, 6, ..., fact(n)}, size = n + 1
vector<ll> factvec(ll n, ll mod) {
  vector<ll> v(n + 1);
  v[0] = 1;
  fr(i, 1, n + 1) { v[i] = (i * v[i - 1]) % mod; }
  return v;
}

ll nPkMod(ll n, ll k, ll mod, const vector<ll> &factvec) {
  ll tmp = 1;
  tmp = (tmp * factvec[n]) % mod;
  tmp = (tmp * invmod(factvec[n - k], mod)) % mod;
  return tmp;
}

ll nCkMod(ll n, ll k, ll mod, const vector<ll> &factvec) {
  ll tmp = nPkMod(n, k, mod, factvec);
  return (tmp * invmod(factvec[k], mod)) % mod;
}

unordered_map<ll, ll> cx;

ll cache(ll m, ll nn) {
  ll MOD = (ll)998244353;
  if (cx.count(nn) > 0)
    return cx[nn] % MOD;
  if (nn <= 1)
    return m % MOD;
  ll base = m;

  cx[nn] = cache(m, nn - 1) * (m - 1) % MOD;
  return cx[nn];
}

int main() {
  ;
  ;
  ll n, m, k;
  cin >> n >> m >> k;
  // 1->n,n-1
  // 2->n,n-1,n-2
  // k->n,n-1,...,n-k (k+1)
  ll sum = 0;
  ll MOD = (ll)998244353;
  auto fvec = factvec(n, MOD);
  rep(i, k + 1) {
    ll nn = n - i;

    auto base = cache(m, nn);
    // i はペア数
    sum = (sum + base * (nCkMod(n - 1, i, MOD, fvec) % MOD)) % MOD;
  }

  pr(sum);
}


Error: TLE

Fixed submission:

#include <algorithm>
#include <cmath>
#include <cstdio>
#include <deque>
#include <functional>
#include <iomanip>
#include <iostream>
#include <map>
#include <numeric>
#include <queue>
#include <random>
#include <set>
#include <stack>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <vector>
#define ll long long
#define fr(i, k, N) for (decltype(N + k) i = k; i < N; i++)
#define rep(i, N) for (decltype(N) i = 0; i < N; i++)
#define ALL(v) v.begin(), v.end()
using namespace std;
template <class X> void pr(X test) { cout << test << endl; }
template <class X> void prr(X test) {
  for (auto it : test)
    cout << it << " ";
  cout << endl;
}
template <class X> void prrr(X test) {
  int f = 0;
  for (auto it : test)
    cout << (f++ != 0 ? " " : "") << it;
  cout << endl;
}
template <typename X> istream &operator>>(istream &i, vector<X> &v) {
  for (X &x : v)
    i >> x;
  return i;
}

class UF { // union find from https://github.com/kartikkukreja/blog-codes (MIT)
           // [usage]
           // auto uf = new UF(n);
           // uf->merge(l, r); // 0 <= l, r <= n-1
           // uf->sz[uf->find(l)]; // size of l's union
           // ...
public:
  int *id; // inner id
  int cnt; // number of all vartices
  int *sz; // sz[i] means size of i's set
  UF(int N) {
    cnt = N;
    id = new int[N];
    sz = new int[N];
    for (int i = 0; i < N; i++) {
      id[i] = i;
      sz[i] = 1;
    }
  }
  ~UF() {
    delete[] id;
    delete[] sz;
  }
  int find(int p) {
    int root = p;
    while (root != id[root])
      root = id[root];
    while (p != root) {
      int newp = id[p];
      id[p] = root;
      p = newp;
    }
    return root;
  }
  void merge(int x, int y) {
    int i = find(x);
    int j = find(y);
    if (i == j)
      return;

    if (sz[i] < sz[j]) {
      id[i] = j;
      sz[j] += sz[i];
    } else {
      id[j] = i;
      sz[i] += sz[j];
    }
    cnt--;
  }
  bool connected(int x, int y) { return find(x) == find(y); }
  int count() { return cnt; }
};

// return 約数unordered_set 1 <= x <= n
auto divv(ll n) {
  unordered_set<decltype(n)> v;
  for (decltype(n) i = 1; i * i <= n; i++) {
    if (n % i == 0) {
      v.insert(i);
      if (i * i != n)
        v.insert(n / i);
    }
  }
  return v;
};

// check prime (and n>=2) or not
auto isPrime(ll n) {
  if (n <= 1)
    return false;
  for (decltype(n) i = 2; i <= sqrt(n); ++i) {
    if (n % i == 0) {
      return false;
    }
  }
  return true;
}

// 最小公倍数
auto lcm_(ll a, ll b) {
  if (a < b)
    swap(a, b);
  auto r = a % b;
  auto x = a * b;
  while (r != 0) {
    a = b;
    b = r;
    r = a % b;
  }
  return x / b;
}

// 最大公約数
auto gcd_(ll a, ll b) { return (a * b) / lcm_(a, b); };

// factorial n mod prm
ll factmod(ll a, ll prm) {
  return a == 0 ? 1 : (a * factmod(a - 1, prm)) % prm;
};

// 1/x mod prm
auto invmod = [](auto x, auto prm) {
  decltype(x) r = 1;
  if (x == 1)
    return r;
  for (auto ex = prm - 2; ex; ex >>= 1) {
    if (ex & 1)
      r = (r * x) % prm;
    x = (x * x) % prm;
  }
  return r;
};

// {{2, 3}, {3, 1}, {5, 2}, ...}
unordered_map<ll, ll> prime_factor(ll n) {
  unordered_map<ll, ll> ret;
  for (ll i = 2; i * i <= n; i++) {
    while (n % i == 0) {
      ret[i]++;
      n /= i;
    }
  }
  if (n != 1)
    ret[n] = 1;
  return ret;
}

ll modpow(ll base, ll exp, ll mod) {
  base %= mod;
  ll result = 1;
  while (exp > 0) {
    if (exp & 1)
      result = (result * base) % mod;
    base = (base * base) % mod;
    exp >>= 1;
  }
  return result;
}

// get {1, 1, 2, 6, ..., fact(n)}, size = n + 1
vector<ll> factvec(ll n, ll mod) {
  vector<ll> v(n + 1);
  v[0] = 1;
  fr(i, 1, n + 1) { v[i] = (i * v[i - 1]) % mod; }
  return v;
}

ll nPkMod(ll n, ll k, ll mod, const vector<ll> &factvec) {
  ll tmp = 1;
  tmp = (tmp * factvec[n]) % mod;
  tmp = (tmp * invmod(factvec[n - k], mod)) % mod;
  return tmp;
}

ll nCkMod(ll n, ll k, ll mod, const vector<ll> &factvec) {
  ll tmp = nPkMod(n, k, mod, factvec);
  return (tmp * invmod(factvec[k], mod)) % mod;
}

int main() {
  ;
  ;
  ll n, m, k;
  cin >> n >> m >> k;
  // 1->n,n-1
  // 2->n,n-1,n-2
  // k->n,n-1,...,n-k (k+1)
  ll sum = 0;
  ll MOD = (ll)998244353;
  auto fvec = factvec(n, MOD);
  rep(i, k + 1) {
    ll nn = n - i;
    ll base = m;
    rep(j, nn - 1) {
      base *= (m - 1);
      base %= MOD;
      ;
    };
    // i はペア数
    sum = (sum + base * (nCkMod(n - 1, i, MOD, fvec) % MOD)) % MOD;
  }

  pr(sum);
}