Problem p02685 (C++)

Buggy submission (status: Runtime Error):
#include <bits/stdc++.h>
#define endl "\n"
using namespace std;
typedef long long int lli;
const lli M = 998244353, N = 1e5 + 5;
lli fact[N];

lli binpow(lli a, lli b) {
  if (b == 0)
    return 1;
  lli res = binpow(a, b / 2);
  res = (res * res) % M;

  if (b % 2)
    return (res * a) % M;
  return res;
}

lli inv(lli a) { return binpow(a, M - 2); }

lli add(lli a, lli b) {
  lli res = a + b;
  if (res >= M)
    return res - M;
  return res;
}

lli mul(lli a, lli b) {
  lli res = a * b;
  if (res >= M)
    return res % M;
  return res;
}

lli sub(lli a, lli b) { return (a % M - b % M + M) % M; }

void pre() {
  fact[0] = 1;
  for (int i = 1; i < N; i++)
    fact[i] = mul(fact[i - 1], i);
}

lli ncr(lli n, lli r) { return mul(fact[n], inv(mul(fact[n - r], fact[r]))); }

/*c(n-1,k1) part is straight forward: find k1 points to split the whole row.

In each of k1 + 1 segments, colors must be different. The 2 blocks at either
side of each splitting point have same color.

m is for the first block of first segment, it can choose from all m colors. for
the rest blocks of first segment, each one can only choose from m-1 colors since
its color must be different from adjacent blocks.

for the first block of other segment, it has only 1 choice: the same as the last
one of previous segment. the rest blocks of other segment are like that in the
first segment, each has m-1 choices.*/

int main() {
  ios::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);
  pre();
  lli n, m, k;
  cin >> n >> m >> k;
  lli i, cur, ans = 0;

  for (i = 0; i <= k; i++) {
    cur = ncr(n - 1, i);
    cur = mul(cur, mul(m, binpow(m - 1, n - 1 - i)));
    ans = add(ans, cur);
  }
  cout << ans;
  return 0;
}

Fixed submission:
#include <bits/stdc++.h>
#define endl "\n"
using namespace std;
typedef long long int lli;
const lli M = 998244353, N = 2e5 + 5;
lli fact[N];

lli binpow(lli a, lli b) {
  if (b == 0)
    return 1;
  lli res = binpow(a, b / 2);
  res = (res * res) % M;

  if (b % 2)
    return (res * a) % M;
  return res;
}

lli inv(lli a) { return binpow(a, M - 2); }

lli add(lli a, lli b) {
  lli res = a + b;
  if (res >= M)
    return res - M;
  return res;
}

lli mul(lli a, lli b) {
  lli res = a * b;
  if (res >= M)
    return res % M;
  return res;
}

lli sub(lli a, lli b) { return (a % M - b % M + M) % M; }

void pre() {
  fact[0] = 1;
  for (int i = 1; i < N; i++)
    fact[i] = mul(fact[i - 1], i);
}

lli ncr(lli n, lli r) { return mul(fact[n], inv(mul(fact[n - r], fact[r]))); }

/*c(n-1,k1) part is straight forward: find k1 points to split the whole row.

In each of k1 + 1 segments, colors must be different. The 2 blocks at either
side of each splitting point have same color.

m is for the first block of first segment, it can choose from all m colors. for
the rest blocks of first segment, each one can only choose from m-1 colors since
its color must be different from adjacent blocks.

for the first block of other segment, it has only 1 choice: the same as the last
one of previous segment. the rest blocks of other segment are like that in the
first segment, each has m-1 choices.*/

int main() {
  ios::sync_with_stdio(0);
  cin.tie(0);
  cout.tie(0);
  pre();
  lli n, m, k;
  cin >> n >> m >> k;
  lli i, cur, ans = 0;

  for (i = 0; i <= k; i++) {
    cur = ncr(n - 1, i);
    cur = mul(cur, mul(m, binpow(m - 1, n - 1 - i)));
    ans = add(ans, cur);
  }
  cout << ans;
  return 0;
}