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

/**
 *    author:  math2do
 *    created: 12.05.2020 22:43:05 IST
 **/

#include <bits/stdc++.h>

using namespace std;

template <typename A, typename B> string to_string(pair<A, B> p);

template <typename A, typename B, typename C>
string to_string(tuple<A, B, C> p);

template <typename A, typename B, typename C, typename D>
string to_string(tuple<A, B, C, D> p);

string to_string(const string &s) { return '"' + s + '"'; }

string to_string(const char *s) { return to_string((string)s); }

string to_string(bool b) { return (b ? "true" : "false"); }

string to_string(const char &ch) { return "'" + string(1, ch) + "'"; }

string to_string(vector<bool> v) {
  bool first = true;
  string res = "{";
  for (int i = 0; i < static_cast<int>(v.size()); i++) {
    if (!first) {
      res += ", ";
    }
    first = false;
    res += to_string(v[i]);
  }
  res += "}";
  return res;
}

template <size_t N> string to_string(bitset<N> v) {
  string res = "";
  for (size_t i = 0; i < N; i++) {
    res += static_cast<char>('0' + v[i]);
  }
  reverse(res.begin(), res.end());
  return res;
}

template <typename A> string to_string(A v) {
  bool first = true;
  string res = "{";
  for (const auto &x : v) {
    if (!first) {
      res += ", ";
    }
    first = false;
    res += to_string(x);
  }
  res += "}";
  return res;
}

template <typename A, typename B> string to_string(pair<A, B> p) {
  return "(" + to_string(p.first) + ", " + to_string(p.second) + ")";
}

template <typename A, typename B, typename C>
string to_string(tuple<A, B, C> p) {
  return "(" + to_string(get<0>(p)) + ", " + to_string(get<1>(p)) + ", " +
         to_string(get<2>(p)) + ")";
}

template <typename A, typename B, typename C, typename D>
string to_string(tuple<A, B, C, D> p) {
  return "(" + to_string(get<0>(p)) + ", " + to_string(get<1>(p)) + ", " +
         to_string(get<2>(p)) + ", " + to_string(get<3>(p)) + ")";
}

void debug_out() { cerr << endl; }

template <typename Head, typename... Tail> void debug_out(Head H, Tail... T) {
  cerr << " " << to_string(H);
  debug_out(T...);
}

#ifdef LOCAL
#define debug(...) cerr << "[" << #__VA_ARGS__ << "]:", debug_out(__VA_ARGS__)
#else
#define debug(...) 42
#endif

template <typename T> T power(T a, long long b) {
  T r = 1;
  while (b > 0) {
    if (b & 1) {
      r *= a;
    }
    a *= a;
    b >>= 1;
  }
  return r;
}

template <typename T> T inverse(T a, T m) {
  a %= m;
  if (a < 0) {
    a += m;
  }
  T b = m, u = 0, v = 1;
  while (a) {
    T t = b / a;
    b -= a * t;
    swap(a, b);
    u -= v * t;
    swap(u, v);
  }
  if (u < 0) {
    u += m;
  }
  return u;
}

template <int _P> struct modnum {
  static constexpr int P = _P;

private:
  int v;

public:
  modnum() : v(0) {}
  modnum(long long _v) {
    v = (int)(_v % P);
    if (v < 0) {
      v += P;
    }
  }
  explicit operator int() const { return v; }
  bool operator==(const modnum &o) const { return v == o.v; }
  bool operator!=(const modnum &o) const { return v != o.v; }
  modnum inverse() const { return modnum(::inverse(v, P)); }
  modnum operator-() const { return modnum(v ? P - v : 0); }
  modnum operator+() const { return *this; }
  modnum &operator++() {
    v++;
    if (v == P) {
      v = 0;
    }
    return *this;
  }
  modnum &operator--() {
    if (v == 0) {
      v = P;
    }
    v--;
    return *this;
  }
  modnum operator++(int) {
    modnum r = *this;
    ++*this;
    return r;
  }
  modnum operator--(int) {
    modnum r = *this;
    --*this;
    return r;
  }
  modnum &operator+=(const modnum &o) {
    v += o.v;
    if (v >= P) {
      v -= P;
    }
    return *this;
  }
  modnum operator+(const modnum &o) const { return modnum(*this) += o; }
  modnum &operator-=(const modnum &o) {
    v -= o.v;
    if (v < 0) {
      v += P;
    }
    return *this;
  }
  modnum operator-(const modnum &o) const { return modnum(*this) -= o; }
  modnum &operator*=(const modnum &o) {
    v = (int)((long long)v * o.v % P);
    return *this;
  }
  modnum operator*(const modnum &o) const { return modnum(*this) *= o; }
  modnum &operator/=(const modnum &o) { return *this *= o.inverse(); }
  modnum operator/(const modnum &o) const { return modnum(*this) /= o; }
};

template <int _P> string to_string(const modnum<_P> &n) {
  return to_string((int)n);
}

template <int _P> ostream &operator<<(ostream &out, const modnum<_P> &n) {
  return out << int(n);
}

template <int _P> istream &operator>>(istream &in, modnum<_P> &n) {
  long long _v;
  in >> _v;
  n = modnum<_P>(_v);
  return in;
}

const int mod = (int)(998244353);
using mint = modnum<mod>;

template <typename T> class combination {
public:
  vector<T> fact;
  vector<T> inv_fact;
  combination(int N) : fact(N + 1), inv_fact(N + 1) {
    fact[0] = 1;
    for (int i = 1; i <= N; i++) {
      fact[i] = fact[i - 1] * i;
    }
    inv_fact[N] = T(1) / fact[N];
    for (int i = N - 1; i >= 0; i--) {
      inv_fact[i] = inv_fact[i + 1] * (i + 1);
    }
  }
  T P(int n, int r) const {
    if (r < 0 || n < r)
      return 0;
    return fact[n] * inv_fact[n - r];
  }

  T C(int n, int r) const {
    if (r < 0 || n < r)
      return 0;
    return fact[n] * inv_fact[r] * inv_fact[n - r];
  }
};

const int N = (int)(2e5 + 1);
combination<mint> co(N);
inline mint C(int n, int r) { return co.C(n, r); }
inline mint P(int n, int r) { return co.P(n, r); }

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(0);
  int n, m, k;
  cin >> n >> m >> k;
  mint ans = 0;
  mint col = m;
  for (int i = n - 1; i >= 0; i--) {
    mint now = col;
    now *= C(n - 1, i);
    if (i <= k)
      ans += now;
    col *= (m - 1);
  }
  cout << ans << endl;
  return 0;
}


Error: TLE

Fixed submission:

/**
 *    author:  math2do
 *    created: 12.05.2020 22:43:05 IST
 **/

#include <bits/stdc++.h>

using namespace std;

template <typename A, typename B> string to_string(pair<A, B> p);

template <typename A, typename B, typename C>
string to_string(tuple<A, B, C> p);

template <typename A, typename B, typename C, typename D>
string to_string(tuple<A, B, C, D> p);

string to_string(const string &s) { return '"' + s + '"'; }

string to_string(const char *s) { return to_string((string)s); }

string to_string(bool b) { return (b ? "true" : "false"); }

string to_string(const char &ch) { return "'" + string(1, ch) + "'"; }

string to_string(vector<bool> v) {
  bool first = true;
  string res = "{";
  for (int i = 0; i < static_cast<int>(v.size()); i++) {
    if (!first) {
      res += ", ";
    }
    first = false;
    res += to_string(v[i]);
  }
  res += "}";
  return res;
}

template <size_t N> string to_string(bitset<N> v) {
  string res = "";
  for (size_t i = 0; i < N; i++) {
    res += static_cast<char>('0' + v[i]);
  }
  reverse(res.begin(), res.end());
  return res;
}

template <typename A> string to_string(A v) {
  bool first = true;
  string res = "{";
  for (const auto &x : v) {
    if (!first) {
      res += ", ";
    }
    first = false;
    res += to_string(x);
  }
  res += "}";
  return res;
}

template <typename A, typename B> string to_string(pair<A, B> p) {
  return "(" + to_string(p.first) + ", " + to_string(p.second) + ")";
}

template <typename A, typename B, typename C>
string to_string(tuple<A, B, C> p) {
  return "(" + to_string(get<0>(p)) + ", " + to_string(get<1>(p)) + ", " +
         to_string(get<2>(p)) + ")";
}

template <typename A, typename B, typename C, typename D>
string to_string(tuple<A, B, C, D> p) {
  return "(" + to_string(get<0>(p)) + ", " + to_string(get<1>(p)) + ", " +
         to_string(get<2>(p)) + ", " + to_string(get<3>(p)) + ")";
}

void debug_out() { cerr << endl; }

template <typename Head, typename... Tail> void debug_out(Head H, Tail... T) {
  cerr << " " << to_string(H);
  debug_out(T...);
}

#ifdef LOCAL
#define debug(...) cerr << "[" << #__VA_ARGS__ << "]:", debug_out(__VA_ARGS__)
#else
#define debug(...) 42
#endif

template <typename T> T power(T a, long long b) {
  T r = 1;
  while (b > 0) {
    if (b & 1) {
      r *= a;
    }
    a *= a;
    b >>= 1;
  }
  return r;
}

template <typename T> T inverse(T a, T m) {
  a %= m;
  if (a < 0) {
    a += m;
  }
  T b = m, u = 0, v = 1;
  while (a) {
    T t = b / a;
    b -= a * t;
    swap(a, b);
    u -= v * t;
    swap(u, v);
  }
  if (u < 0) {
    u += m;
  }
  return u;
}

template <int _P> struct modnum {
  static constexpr int P = _P;

private:
  int v;

public:
  modnum() : v(0) {}
  modnum(long long _v) {
    v = (int)(_v % P);
    if (v < 0) {
      v += P;
    }
  }
  explicit operator int() const { return v; }
  bool operator==(const modnum &o) const { return v == o.v; }
  bool operator!=(const modnum &o) const { return v != o.v; }
  modnum inverse() const { return modnum(::inverse(v, P)); }
  modnum operator-() const { return modnum(v ? P - v : 0); }
  modnum operator+() const { return *this; }
  modnum &operator++() {
    v++;
    if (v == P) {
      v = 0;
    }
    return *this;
  }
  modnum &operator--() {
    if (v == 0) {
      v = P;
    }
    v--;
    return *this;
  }
  modnum operator++(int) {
    modnum r = *this;
    ++*this;
    return r;
  }
  modnum operator--(int) {
    modnum r = *this;
    --*this;
    return r;
  }
  modnum &operator+=(const modnum &o) {
    v += o.v;
    if (v >= P) {
      v -= P;
    }
    return *this;
  }
  modnum operator+(const modnum &o) const { return modnum(*this) += o; }
  modnum &operator-=(const modnum &o) {
    v -= o.v;
    if (v < 0) {
      v += P;
    }
    return *this;
  }
  modnum operator-(const modnum &o) const { return modnum(*this) -= o; }
  modnum &operator*=(const modnum &o) {
    v = (int)((long long)v * o.v % P);
    return *this;
  }
  modnum operator*(const modnum &o) const { return modnum(*this) *= o; }
  modnum &operator/=(const modnum &o) { return *this *= o.inverse(); }
  modnum operator/(const modnum &o) const { return modnum(*this) /= o; }
};

template <int _P> string to_string(const modnum<_P> &n) {
  return to_string((int)n);
}

template <int _P> ostream &operator<<(ostream &out, const modnum<_P> &n) {
  return out << int(n);
}

template <int _P> istream &operator>>(istream &in, modnum<_P> &n) {
  long long _v;
  in >> _v;
  n = modnum<_P>(_v);
  return in;
}

const int mod = (int)(998244353);
using mint = modnum<mod>;

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(0);
  int n, m, k;
  cin >> n >> m >> k;
  function<mint(int, int, int)> get = [&](int last, int i, int len) {
    if (i == n) {
      return mint(1);
    }
    mint ans = 0;
    for (int c = 0; c < m; c++) {
      if (c == last) {
        if (len < k + 1) {
          ans += get(c, i + 1, len + 1);
        }
      } else {
        ans += get(c, i + 1, 1);
      }
    }
    return ans;
  };
  mint ans = 0;
  for (int c = 0; c < m; c++) {
    ans += get(c, 1, 1);
  }
  cout << ans << endl;
  return 0;
}