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

#include <cmath>
#include <iostream>
#include <map>
#include <set>
#include <vector>

using namespace std;
typedef long long int ll;
typedef unsigned long long int ull;

#define MOD 998244353LL

ll N, M, K;
ll pat(ll, ll);
ll mpow(ll, ll);
ll C(int, int);

map<pair<ll, ll>, ll> mp;
map<ll, ll> mf;

int main() {
  ll ans;

  cin >> N >> M >> K;

  ans = pat(N, K);

  cout << ans << endl;
  return 0;
}
ll pat(ll n, ll k) {
  ll ret = 0;
  ll i;
  for (i = 0; i <= k; i++) {
    ret += M * mpow(M - 1, n - i - 1) % MOD * C(n - 1, i) % MOD;
    ret %= MOD;
  }
  return ret;
}
ll _pat(ll n, ll k) {
  ll ret;

  // printf("%*spat(%lld, %lld)\\n", (int)(N-n), "", n, k);
  if (n == 1) {
    // printf("%*spat(%lld, %lld)=1\\n", (int)(N-n), "", n, k);
    return M;
  }
  auto it = mp.find(make_pair(n, min(n, k + 1)));
  if (it != mp.end()) {
    return it->second;
  }

  if (n <= k + 1) { // 隣り合うブロックの色が同じ数の制限がない
    ret = mpow(M, n);
    mp.emplace(make_pair(n, min(n, k + 1)), ret);
    return ret;
  }
  if (k == 0) { // 隣り合うブロックはすべて異なる色の場合
    ret = M * mpow(M - 1, n - 1);
    ret %= MOD;
    mp.emplace(make_pair(n, min(n, 1LL)), ret);
    return ret;
  }

  if (k == 1) { // 必ずn>=2
    ret = M * mpow(M - 1, n - 2) % MOD * ((n - 1) + (M - 1)) % MOD;
    // printf("%lld pat(%lld, %lld)\\n", ret, n, k);
    mp.emplace(make_pair(n, min(n, k + 1)), ret);
    return ret;
  }

  // 前のブロックと異なる色で始まるものの個数
  ret = (M - 1) * pat(n - 1, k) % MOD;
  // 前のブロックと同じ色で始まるものの個数
  if (k > 0) {
    ret += pat(n - 1, k - 1);
    ret %= MOD;
  }
  mp.emplace(make_pair(n, min(n, k + 1)), ret);
  // printf("%*spat(%lld, %lld)=%lld\\n", (int)(N-n), "", n, k, ret);
  // printf("%lld pat(%lld, %lld)\\n", ret, n, k);
  return ret;
}
ll mpow(ll m, ll n) {
  ll ret;
  if (n == 0)
    return 1;
  if (n == 1)
    return m;
  auto it = mp.find(make_pair(m, n));
  if (it != mp.end())
    return it->second;
  ret = m * mpow(m, n - 1) % MOD;
  mp.emplace(make_pair(m, n), ret);
  // printf("mpow(%lld, %lld) = %lld\\n", m, n, ret);
  return ret;
}
ll _mpow(ll m, ll n) {
  ll i;
  ll ret = m;
  if (n == 0)
    return 1;
  for (i = 1; i < n; i++) {
    ret *= m;
    ret %= MOD;
  }
  // printf("mpow(%lld, %lld)\\n", m, n);
  return ret;
}
int xgcd(int, int, int &, int &);
int gcd(int, int);
int modinv(int, int);
ll modInverse(ll, ll);
ll power(ll, ll, ll);

ll pow2(int);

/*
  nCkを求める
  MOD(1000000007LL)で割った余りを算出
  */
ll fact(int n) {
  ll ret = 1;
  if (n == 0 || n == 1)
    return 1;
  auto it = mf.find(n);
  if (it != mf.end())
    return it->second;
  ret = n * fact(n - 1) % MOD;
  mf.emplace(n, ret);
  // printf("fact(%d) = %lld\\n", n, ret);
  return ret;
}
ll C(int n, int k) {
  ll f1, f2, f3;
  ll ret;

  f1 = fact(n);
  f2 = fact(k);
  f3 = fact(n - k);
  ret = f1;
  ret *= modinv(f2, MOD);
  ret %= MOD;
  ret *= modinv(f3, MOD);
  ret %= MOD;
  return ret;
}
ll _C(int n, int k) {
  ll ret = 1;
  int i;
  int a, b;

  a = k;
  b = n - k;
  if (a < b)
    swap(a, b);

  for (i = n; i > a; i--) {
    ret *= i;
    ret %= MOD;
  }
  for (i = 2; i <= b; i++) {
    ret *= modinv(i, MOD);
    ret %= MOD;
  }
  // printf("C(%d, %d)\\n", n, k);
  return ret;
}

// 素数判定

// 除算
//  素因数分解された値の除算
map<int, int> vdiv(map<int, int> m, map<int, int> n) {
  map<int, int> ret;

  for (auto itm = m.begin(); itm != m.end(); itm++) {
    auto itn = n.find(itm->first);
    if (itn == n.end()) {
      ret[itm->first] = itm->second;
    } else {
      ret[itm->first] = itm->second - itn->second;
      n.erase(itm->first);
    }
  }
  for (auto itn = n.begin(); itn != n.end(); itn++) {
    ret[itn->first] = -itn->second;
  }
  return ret;
}
// modの剰余として、素因数分解したものをint型に変換する
int toInt(map<int, int> f, int m) {
  ll lret = 1;

  for (auto itr = f.begin(); itr != f.end(); itr++) {
    for (int j = 0; j < itr->second; j++) {
      lret *= itr->first;
      lret %= m;
    }
  }
  return lret;
}

// 最大公約数 ユークリッド互除法
int gcd(int m, int n) {
  if (m < n)
    swap(m, n);
  while (m % n) {
    m %= n;
    swap(m, n);
  }
  return n;
}
// 最大公約数 拡張ユークリッド互除法
int xgcd(int a, int b, int &x0, int &y0) {
  x0 = 1, y0 = 0;
  int x1 = 0, y1 = 1;
  int q;
  int at, x0t, y0t;
  while (b != 0) {
    q = a / b;
    at = b;
    b = a % b;
    a = at;
    x0t = x1;
    x1 = x0 - q * x1;
    x0 = x0t;
    y0t = y1;
    y1 = y0 - q * y1;
    y0 = y0t;
  }
  return a;
}
ll xgcd(ll a, ll b, ll &x0, ll &y0) {
  x0 = 1, y0 = 0;
  ll x1 = 0, y1 = 1;
  ll q;
  ll at, x0t, y0t;
  while (b != 0) {
    q = a / b;
    at = b;
    b = a % b;
    a = at;
    x0t = x1;
    x1 = x0 - q * x1;
    x0 = x0t;
    y0t = y1;
    y1 = y0 - q * y1;
    y0 = y0t;
  }
  return a;
}
// モジュラー逆数
//((ab mod n)(b^-1 mod n)) mod n = a mod n となる aを求めるのに使う
// ab mod n と b mod nがわかっているときに a mod n を算出するため
// modInverseは実績ありmodinvは不明
int modinv(int a, int m) {
  int g, x, y;
  g = xgcd(a, m, x, y);
  if (g != 1)
    return -1;
  x %= m;
  if (x < 0)
    x += m;
  return x;
}
ll modinv(ll a, ll m) {
  ll g, x, y;
  g = xgcd(a, m, x, y);
  if (g != 1)
    return -1;
  x %= m;
  if (x < 0)
    x += m;
  return x;
}
ll power(ll x, ll y, ll p) {
  ll res = 1;
  x = x % p;
  while (y > 0) {
    if (y & 1)
      res = (res * x) % p;
    y = y >> 1;
    x = (x * x) % p;
  }
  return res;
}
ll modInverse(ll a, ll m) {
  ll ans = power(a, m - 2, m);
  return ans;
}

ll pow2(int n) {
  ll ret = 1;
  ll r2 = 2;
  int i;

  for (i = n; i > 0; i /= 2) {
    if (i % 2)
      ret *= r2;
    ret %= MOD;
    r2 *= r2;
    r2 %= MOD;
  }
  return ret;
}

/*
  ブロックの塗り方は M^N通り
  隣り合うブロックの色が同じ個所
  0個:  M*(M-1)^(N-1)
    最初のブロックの色はM通り2番目以降は前の色と異なるのでM-1通り
  1個:  (M*(M-1)^(N-2))*(N-1)
    同色の個所はN-1通り、同色以外のブロックはM-1色のN-2個分
  2個:  M*(M-1)^(N-3)*(N-1)*(N-2)/2
    隣り合うブロックが同色のものを1まとめと考えると、
    N-2個のグループに分けられる
    同色の個所はN-1C2=(N-1)*(N-2)/2通り
  k個:  M*(M-1)^(N-k-1)*((N-1)Ck)

  3 3 1
  112 113 121 122 123 131 132 133
  */


Error: TLE

Fixed submission:

#include <cmath>
#include <iostream>
#include <map>
#include <set>
#include <vector>

using namespace std;
typedef long long int ll;
typedef unsigned long long int ull;

#define MOD 998244353LL

ll N, M, K;
ll pat(ll, ll);
ll mpow(ll, ll);
ll C(int, int);

map<pair<ll, ll>, ll> mp, mc;

int main() {
  ll ans;

  cin >> N >> M >> K;

  ans = pat(N, K);

  cout << ans << endl;
  return 0;
}
ll pat(ll n, ll k) {
  ll ret = 0;
  ll i;
  for (i = 0; i <= k; i++) {
    ret += M * mpow(M - 1, n - i - 1) % MOD * C(n - 1, i) % MOD;
    ret %= MOD;
  }
  return ret;
}
ll _pat(ll n, ll k) {
  ll ret;

  // printf("%*spat(%lld, %lld)\\n", (int)(N-n), "", n, k);
  if (n == 1) {
    // printf("%*spat(%lld, %lld)=1\\n", (int)(N-n), "", n, k);
    return M;
  }
  auto it = mp.find(make_pair(n, min(n, k + 1)));
  if (it != mp.end()) {
    return it->second;
  }

  if (n <= k + 1) { // 隣り合うブロックの色が同じ数の制限がない
    ret = mpow(M, n);
    mp.emplace(make_pair(n, min(n, k + 1)), ret);
    return ret;
  }
  if (k == 0) { // 隣り合うブロックはすべて異なる色の場合
    ret = M * mpow(M - 1, n - 1);
    ret %= MOD;
    mp.emplace(make_pair(n, min(n, 1LL)), ret);
    return ret;
  }

  if (k == 1) { // 必ずn>=2
    ret = M * mpow(M - 1, n - 2) % MOD * ((n - 1) + (M - 1)) % MOD;
    // printf("%lld pat(%lld, %lld)\\n", ret, n, k);
    mp.emplace(make_pair(n, min(n, k + 1)), ret);
    return ret;
  }

  // 前のブロックと異なる色で始まるものの個数
  ret = (M - 1) * pat(n - 1, k) % MOD;
  // 前のブロックと同じ色で始まるものの個数
  if (k > 0) {
    ret += pat(n - 1, k - 1);
    ret %= MOD;
  }
  mp.emplace(make_pair(n, min(n, k + 1)), ret);
  // printf("%*spat(%lld, %lld)=%lld\\n", (int)(N-n), "", n, k, ret);
  // printf("%lld pat(%lld, %lld)\\n", ret, n, k);
  return ret;
}
ll mpow(ll m, ll n) {
  ll i;
  ll ret = m;
  if (n == 0)
    return 1;
  for (i = 1; i < n; i++) {
    ret *= m;
    ret %= MOD;
  }
  // printf("mpow(%lld, %lld)\\n", m, n);
  return ret;
}
int xgcd(int, int, int &, int &);
int gcd(int, int);
int modinv(int, int);
ll modInverse(ll, ll);
ll power(ll, ll, ll);

ll pow2(int);

/*
  nCkを求める
  MOD(1000000007LL)で割った余りを算出
  */
ll C(int n, int k) {
  ll ret = 1;
  int i;
  int a, b;

  a = k;
  b = n - k;
  if (a < b)
    swap(a, b);

  for (i = n; i > a; i--) {
    ret *= i;
    ret %= MOD;
  }
  for (i = 2; i <= b; i++) {
    ret *= modinv(i, MOD);
    ret %= MOD;
  }
  // printf("C(%d, %d)\\n", n, k);
  return ret;
}

// 素数判定
vector<int> pr = {2, 3, 5, 7};

bool isprime(int n) {
  int i;
  int prime(int);
  int sz = pr.size();
  int sq = sqrt(n);
  for (i = 1; i < sz; i++) {
    int p = pr[i];
    if (p > sq)
      return true;
    if (n % p == 0)
      return false;
  }
  for (;; i++) {
    int p = prime(i);
    if (p > sq)
      return true;
    if (n % p == 0)
      return false;
  }
}

// 素数生成
int prime(int n) {
  int last = pr.size();
  if (n < last)
    return pr[n];

  for (int p = pr.back() + 2; last <= n; p += 2) {
    if (isprime(p)) {
      last++;
      pr.push_back(p);
    }
  }
  return pr.back();
}

// 素因数分解
// 引数の整数を素因数分解する。
// mapクラスのキーに基数、対応する値に指数が入っている
map<int, int> lc;
map<int, int> fact(int n) {
  map<int, int> ret;
  int base, ix;

  int i;
  for (i = 0; n > 1; i++) {
    base = prime(i);
    ix = 0;
    while (n % base == 0) {
      ix++;
      n /= base;
    }
    if (ix > 0) {
      ret[base] = ix;
      auto itr = lc.find(base);
      if (itr == lc.end()) {
        lc[base] = ix;
      } else {
        lc[base] = max(ix, lc[base]);
      }
    } else {
      if (base * base >= n) {
        ret[n] = 1;
        auto itr = lc.find(n);
        if (itr == lc.end()) {
          lc[n] = 1;
        } else {
          lc[n] = max(1, lc[n]);
        }
        break;
      }
    }
  }
  return ret;
}

// 最大公約数
// 素因数分解された値での最大公約数算出
map<int, int> lcm(map<int, int> m, map<int, int> n) {
  map<int, int> ret;

  for (auto itm = m.begin(); itm != m.end(); itm++) {
    auto itn = n.find(itm->first);
    if (itn == n.end()) {
      ret[itm->first] = itm->second;
    } else {
      ret[itm->first] = max(itm->second, itn->second);
      n.erase(itm->first);
    }
  }
  for (auto itn = n.begin(); itn != n.end(); itn++) {
    ret[itn->first] = itn->second;
  }

  return ret;
}

// 除算
//  素因数分解された値の除算
map<int, int> vdiv(map<int, int> m, map<int, int> n) {
  map<int, int> ret;

  for (auto itm = m.begin(); itm != m.end(); itm++) {
    auto itn = n.find(itm->first);
    if (itn == n.end()) {
      ret[itm->first] = itm->second;
    } else {
      ret[itm->first] = itm->second - itn->second;
      n.erase(itm->first);
    }
  }
  for (auto itn = n.begin(); itn != n.end(); itn++) {
    ret[itn->first] = -itn->second;
  }
  return ret;
}
// modの剰余として、素因数分解したものをint型に変換する
int toInt(map<int, int> f, int m) {
  ll lret = 1;

  for (auto itr = f.begin(); itr != f.end(); itr++) {
    for (int j = 0; j < itr->second; j++) {
      lret *= itr->first;
      lret %= m;
    }
  }
  return lret;
}

// 最大公約数 ユークリッド互除法
int gcd(int m, int n) {
  if (m < n)
    swap(m, n);
  while (m % n) {
    m %= n;
    swap(m, n);
  }
  return n;
}
// 最大公約数 拡張ユークリッド互除法
int xgcd(int a, int b, int &x0, int &y0) {
  x0 = 1, y0 = 0;
  int x1 = 0, y1 = 1;
  int q;
  int at, x0t, y0t;
  while (b != 0) {
    q = a / b;
    at = b;
    b = a % b;
    a = at;
    x0t = x1;
    x1 = x0 - q * x1;
    x0 = x0t;
    y0t = y1;
    y1 = y0 - q * y1;
    y0 = y0t;
  }
  return a;
}
ll xgcd(ll a, ll b, ll &x0, ll &y0) {
  x0 = 1, y0 = 0;
  ll x1 = 0, y1 = 1;
  ll q;
  ll at, x0t, y0t;
  while (b != 0) {
    q = a / b;
    at = b;
    b = a % b;
    a = at;
    x0t = x1;
    x1 = x0 - q * x1;
    x0 = x0t;
    y0t = y1;
    y1 = y0 - q * y1;
    y0 = y0t;
  }
  return a;
}
// モジュラー逆数
//((ab mod n)(b^-1 mod n)) mod n = a mod n となる aを求めるのに使う
// ab mod n と b mod nがわかっているときに a mod n を算出するため
// modInverseは実績ありmodinvは不明
int modinv(int a, int m) {
  int g, x, y;
  g = xgcd(a, m, x, y);
  if (g != 1)
    return -1;
  x %= m;
  if (x < 0)
    x += m;
  return x;
}
ll modinv(ll a, ll m) {
  ll g, x, y;
  g = xgcd(a, m, x, y);
  if (g != 1)
    return -1;
  x %= m;
  if (x < 0)
    x += m;
  return x;
}
ll power(ll x, ll y, ll p) {
  ll res = 1;
  x = x % p;
  while (y > 0) {
    if (y & 1)
      res = (res * x) % p;
    y = y >> 1;
    x = (x * x) % p;
  }
  return res;
}
ll modInverse(ll a, ll m) {
  ll ans = power(a, m - 2, m);
  return ans;
}

ll pow2(int n) {
  ll ret = 1;
  ll r2 = 2;
  int i;

  for (i = n; i > 0; i /= 2) {
    if (i % 2)
      ret *= r2;
    ret %= MOD;
    r2 *= r2;
    r2 %= MOD;
  }
  return ret;
}

/*
  ブロックの塗り方は M^N通り
  隣り合うブロックの色が同じ個所
  0個:  M*(M-1)^(N-1)
    最初のブロックの色はM通り2番目以降は前の色と異なるのでM-1通り
  1個:  (M*(M-1)^(N-2))*(N-1)
    同色の個所はN-1通り、同色以外のブロックはM-1色のN-2個分
  2個:  M*(M-1)^(N-3)*(N-1)*(N-2)/2
    隣り合うブロックが同色のものを1まとめと考えると、
    N-2個のグループに分けられる
    同色の個所はN-1C2=(N-1)*(N-2)/2通り
  k個:  M*(M-1)^(N-k-1)*((N-1)Ck)

  3 3 1
  112 113 121 122 123 131 132 133
  */