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

// 問題の URL を書いておく
// https://atcoder.jp/contests/abc167/tasks/abc167_e

#include <algorithm>
#include <bitset>
#include <iostream>
#include <map>
#include <numeric>
#include <vector>
using namespace std;

// #define ENABLE_PRINT

template <typename T> T Input() {
  T v;
  cin >> v;
  return v;
}

template <typename T> vector<T> InputVector(int n) {
  vector<T> v(n);
  for (int i = 0; i < n; ++i) {
    cin >> v[i];
  }
  return v;
}

template <typename T> void InputArray(T *begin, int n) {
  for (int i = 0; i < n; ++i) {
    cin >> begin[i];
  }
}

#if defined(ENABLE_PRINT)

#define Print(v)                                                               \\
  do {                                                                         \\
    cout << #v << ": " << v << endl;                                           \\
  } while (0)

#define PrintVec(v)                                                            \\
  do {                                                                         \\
    for (int __i = 0; __i < v.size(); ++__i) {                                 \\
      cout << #v << "[" << __i << "]: " << v[__i] << endl;                     \\
    }                                                                          \\
  } while (0)

#else

#define Print(v) ((void)0)
#define PrintVec(v) ((void)0)

#endif

template <int Mod> struct MintBase {
  int64_t x; // typedef long long int64_t;
  MintBase(int64_t x = 0) : x((x % Mod + Mod) % Mod) {}
  MintBase operator-() const { return MintBase(-x); }
  MintBase &operator+=(const MintBase a) {
    if ((x += a.x) >= Mod)
      x -= Mod;
    return *this;
  }
  MintBase &operator-=(const MintBase a) {
    if ((x += Mod - a.x) >= Mod)
      x -= Mod;
    return *this;
  }
  MintBase &operator*=(const MintBase a) {
    (x *= a.x) %= Mod;
    return *this;
  }
  MintBase operator+(const MintBase a) const { return MintBase(*this) += a; }
  MintBase operator-(const MintBase a) const { return MintBase(*this) -= a; }
  MintBase operator*(const MintBase a) const { return MintBase(*this) *= a; }
  MintBase pow(int64_t t) const {
    if (!t)
      return 1;
    MintBase a = pow(t >> 1);
    a *= a;
    if (t & 1)
      a *= *this;
    return a;
  }

  // for prime Mod
  MintBase inv() const { return pow(Mod - 2); }
  MintBase &operator/=(const MintBase a) { return *this *= a.inv(); }
  MintBase operator/(const MintBase a) const { return MintBase(*this) /= a; }
};

// const int Mod = 1000000007;
// const int Mod = 998244353;

using mint = MintBase<998244353>;

istream &operator>>(istream &is, const mint &a) { return is >> a.x; }
ostream &operator<<(ostream &os, const mint &a) { return os << a.x; }

mint factorial(const mint &x) {
  mint ans = 1;
  for (mint i = x.x; i.x != 0; i -= 1) {
    ans *= i;
  }
  return ans;
}

mint factorial(const mint &x, const mint &n) {
  mint ans = 1;
  mint m = x;
  for (int64_t i = 0; i < n.x; ++i) {
    ans *= m;
    m -= 1;
  }
  return ans;
}

mint Combination(int64_t x, int64_t y) {
  return factorial(x, y) / factorial(y);
}

int main(int argc, const char *argv[]) {
  int n, m, k;
  cin >> n >> m >> k;

  vector<mint> mul(n);
  mul[0] = m;
  for (int i = 1; i < n; ++i) {
    mul[i] = mul[i - 1] * (m - 1);
  }

  mint ans = 0;
  mint upper = 1;
  mint lower = 1;
  for (int i = 0; i <= k; ++i) {
    mint c = 1;
    if (i > 0) {
      upper *= (n - i);
      lower *= i;
      c = upper / lower;
    }
    auto paint = mul[n - 1 - i];
    ans += c * paint;
  }

  cout << ans << endl;
}


Error: TLE

Fixed submission:

// 問題の URL を書いておく
// https://atcoder.jp/contests/abc167/tasks/abc167_e

#include <algorithm>
#include <bitset>
#include <iostream>
#include <map>
#include <numeric>
#include <vector>
using namespace std;

// #define ENABLE_PRINT

template <typename T> T Input() {
  T v;
  cin >> v;
  return v;
}

template <typename T> vector<T> InputVector(int n) {
  vector<T> v(n);
  for (int i = 0; i < n; ++i) {
    cin >> v[i];
  }
  return v;
}

template <typename T> void InputArray(T *begin, int n) {
  for (int i = 0; i < n; ++i) {
    cin >> begin[i];
  }
}

#if defined(ENABLE_PRINT)

#define Print(v)                                                               \\
  do {                                                                         \\
    cout << #v << ": " << v << endl;                                           \\
  } while (0)

#define PrintVec(v)                                                            \\
  do {                                                                         \\
    for (int __i = 0; __i < v.size(); ++__i) {                                 \\
      cout << #v << "[" << __i << "]: " << v[__i] << endl;                     \\
    }                                                                          \\
  } while (0)

#else

#define Print(v) ((void)0)
#define PrintVec(v) ((void)0)

#endif

template <int Mod> struct MintBase {
  int64_t x; // typedef long long int64_t;
  MintBase(int64_t x = 0) : x((x % Mod + Mod) % Mod) {}
  MintBase operator-() const { return MintBase(-x); }
  MintBase &operator+=(const MintBase a) {
    if ((x += a.x) >= Mod)
      x -= Mod;
    return *this;
  }
  MintBase &operator-=(const MintBase a) {
    if ((x += Mod - a.x) >= Mod)
      x -= Mod;
    return *this;
  }
  MintBase &operator*=(const MintBase a) {
    (x *= a.x) %= Mod;
    return *this;
  }
  MintBase operator+(const MintBase a) const { return MintBase(*this) += a; }
  MintBase operator-(const MintBase a) const { return MintBase(*this) -= a; }
  MintBase operator*(const MintBase a) const { return MintBase(*this) *= a; }
  MintBase pow(int64_t t) const {
    if (!t)
      return 1;
    MintBase a = pow(t >> 1);
    a *= a;
    if (t & 1)
      a *= *this;
    return a;
  }

  // for prime Mod
  MintBase inv() const { return pow(Mod - 2); }
  MintBase &operator/=(const MintBase a) { return *this *= a.inv(); }
  MintBase operator/(const MintBase a) const { return MintBase(*this) /= a; }
};

// const int Mod = 1000000007;
// const int Mod = 998244353;

using mint = MintBase<998244353>;

istream &operator>>(istream &is, const mint &a) { return is >> a.x; }
ostream &operator<<(ostream &os, const mint &a) { return os << a.x; }

mint factorial(const mint &x) {
  mint ans = 1;
  for (mint i = x.x; i.x != 0; i -= 1) {
    ans *= i;
  }
  return ans;
}

mint factorial(const mint &x, const mint &n) {
  mint ans = 1;
  mint m = x;
  for (int64_t i = 0; i < n.x; ++i) {
    ans *= m;
    m -= 1;
  }
  return ans;
}

mint Combination(int64_t x, int64_t y) {
  return factorial(x, y) / factorial(y);
}

int main(int argc, const char *argv[]) {
  int n, m, k;
  cin >> n >> m >> k;

  vector<mint> mul(n);
  mul[0] = m;
  for (int i = 1; i < n; ++i) {
    mul[i] = mul[i - 1] * (m - 1);
  }

  mint ans = 0;
  for (int i = 0; i <= k; ++i) {
    auto c = (i == 0 ? 1 : Combination(n - 1, i));
    auto paint = mul[n - 1 - i];
    ans += c * paint;
  }

  cout << ans << endl;
}