CF1528F AmShZ Farm 题解

思路#

首先，原题意要你求：

\sum_{i = 0}^{n} (\binom{n}{i}) i^{k} n^{n - i}

其实比较板子吧。

\sum_{i = 0}^{n} (\binom{n}{i}) i^{k} n^{n - i}

\begin{aligned} = \sum_{i = 0}^{n} (\binom{n}{i}) \sum_{j = 1}^{k} {\binom{k}{j}} (\binom{i}{j}) j! n^{n - i} & (普通幂拆开) \\ = \sum_{j = 1}^{k} {\binom{k}{j}} j! \sum_{i = 0}^{n} (\binom{n}{i}) (\binom{i}{j}) n^{n - i} \\ = \sum_{j = 1}^{k} {\binom{k}{j}} j! \sum_{i = 0}^{n} \frac{n! i!}{i! (n - i)! j! (i - j)!} n^{n - i} \\ = \sum_{j = 1}^{k} {\binom{k}{j}} n! \sum_{i = 0}^{n} \frac{n^{n - i}}{(n - i)! (i - j)!} \\ = \sum_{j = 1}^{k} {\binom{k}{j}} \frac{n!}{(n - j)!} \sum_{i = 0}^{n} (\binom{n - j}{n - i}) n^{n - i} \\ = \sum_{j = 1}^{k} {\binom{k}{j}} \frac{n!}{(n - j)!} (1 + n)^{n - j} & (二项式定理) \\ = \sum_{j = 1}^{k} \frac{1}{j!} \sum_{i = 0}^{j} (\binom{j}{i}) (- 1)^{j - i} i^{k} \frac{n!}{(n - j)!} (1 + n)^{n - j} & (斯特林数拆开) \\ = \sum_{j = 1}^{k} (\binom{n}{j}) \sum_{i = 0}^{j} (\binom{j}{i}) (- 1)^{j - i} i^{k} (1 + n)^{n - j} \\ = \sum_{i = 1}^{k} i^{k} (- 1)^{i} \sum_{j = i}^{k} (\binom{n}{j}) (\binom{j}{i}) (- 1)^{j} (1 + n)^{n - j} \\ = (1 + n)^{n} \sum_{i = 1}^{k} i^{k} (- 1)^{i} \sum_{j = i}^{k} (\binom{n}{j}) (\binom{j}{i}) (\frac{- 1}{1 + n})^{j} \\ = (1 + n)^{n} \sum_{i = 1}^{k} i^{k} (- 1)^{i} \sum_{j = i}^{k} \frac{n! j!}{(n - j)! j! i! (j - i)!} (\frac{- 1}{1 + n})^{j} \\ = (1 + n)^{n} \sum_{i = 1}^{k} i^{k} (- 1)^{i} (\binom{n}{i}) \sum_{j = i}^{k} (\binom{n - i}{j - i}) (\frac{- 1}{1 + n})^{j} \\ = (1 + n)^{n} \sum_{i = 1}^{k} i^{k} (\binom{n}{i}) (\frac{1}{1 + n})^{i} \sum_{j = 0}^{k - i} (\binom{n - i}{j}) (\frac{- 1}{1 + n})^{j} \end{aligned}

设 $f_{i} = \sum_{j = 0}^{k - i} (\binom{n - i}{j}) (\frac{- 1}{1 + n})^{j}$ 。

那么：

\begin{aligned} f_{i} - f_{i + 1} & = \sum_{j = 0}^{k - i} (\binom{n - i}{j}) (\frac{- 1}{1 + n})^{j} - \sum_{j = 0}^{k - i - 1} (\binom{n - i - 1}{j}) (\frac{- 1}{1 + n})^{j} \\ = (\binom{n - i}{k - i}) (\frac{- 1}{1 + n})^{k - i} + \sum_{j = 1}^{k - i - 1} (\binom{n - i - 1}{j - 1}) (\frac{- 1}{1 + n})^{j} \\ = (\binom{n - i}{k - i}) (\frac{- 1}{1 + n})^{k - i} + \sum_{j = 0}^{k - i - 2} (\binom{n - i - 1}{j}) (\frac{- 1}{1 + n})^{j + 1} \\ = (\binom{n - i}{k - i}) (\frac{- 1}{1 + n})^{k - i} + f_{i + 1} \times \frac{- 1}{1 + n} - (\binom{n - i - 1}{k - i - 1}) (\frac{- 1}{1 + n})^{k - i} \\ = (\binom{n - i - 1}{k - i}) (\frac{- 1}{1 + n})^{k - i} + f_{i + 1} \times \frac{- 1}{1 + n} \end{aligned}

所以：

\begin{array}{r} f_{i} = (\binom{n - i - 1}{k - i}) (\frac{- 1}{1 + n})^{k - i} + f_{i + 1} (1 + \frac{- 1}{1 + n}) \end{array}

递推即可。

发现上面所有的东西都可以线性处理。

时间复杂度： $O (k)$ 。

Code#

#include <bits/stdc++.h>
using namespace std;

#define fro(i, x, y) for (int i = (x); i <= (y); i++)
#define pre(i, x, y) for (int i = (x); i >= (y); i--)

typedef long long i64;

const int N = 1e6 + 10;
const int mod = 998244353;

int n, k, ct;
int r, num, sum, ans;
int f[N], iv[N], fc[N], vs[N], pr[N], sm[N];

inline i64 power(i64 x, i64 y) {
  i64 res = 1;
  while (y) {
    if (y & 1) res = res * x % mod;
    x = x * x % mod, y /= 2;
  }
  return res;
}
inline i64 C1(int x, int y) {
  if (x < 0 || y < 0 || x - y < 0) return 0;
  while (r < x) ++r, num = 1ll * num * r % mod;
  return 1ll * num * iv[y] % mod;
}
inline i64 C2(int x, int y) {
  if (x < 0 || y < 0 || x - y < 0) return 0;
  return 1ll * (sum = 1ll * sum * (n - y + 1) % mod) * iv[y] % mod;
}

signed main() {
  ios::sync_with_stdio(0), cin.tie(0);
  cin >> n >> k;
  i64 p1 = mod - power(1 + n, mod - 2);
  i64 p2 = mod - p1;
  fc[0] = 1;
  fro(i, 1, k) fc[i] = 1ll * fc[i - 1] * i % mod;
  iv[k] = power(fc[k], mod - 2);
  pre(i, k, 1) iv[i - 1] = 1ll * iv[i] * i % mod;
  r = n - k, sum = 1, num = n - k;
  if (n > k) {
    f[k] = 1;
    for (int i = k - 1, pw = p1; i >= 0; i--, pw = p1 * pw % mod) {
      f[i] = (f[i + 1] * (p1 + 1) + C1(n - i - 1, k - i) * pw) % mod;
    }
  } else {
    i64 p = mod + 1 - p2; f[n] = 1;
    for (int i = n - 1; i >= 1; i--) {
      f[i] = f[i + 1] * p % mod;
    }
  }
  sm[1] = 1;
  fro(i, 2, k) {
    if (vs[i] == 0) pr[++ct] = i, sm[i] = power(i, k);
    for (int j = 1; j <= ct && pr[j] * i <= k; j++) {
      vs[pr[j] * i] = 1;
      sm[pr[j] * i] = 1ll * sm[pr[j]] * sm[i] % mod;
      if (i % pr[j] == 0) break;
    }
  }
  for (int i = 1, pw = p2; i <= k; i++, pw = p2 * pw % mod) {
    i64 val = sm[i];
    val = val * C2(n, i) % mod;
    val = val * pw % mod;
    ans = (ans + f[i] * val) % mod;
  }
  ans = ans * power(1 + n, n) % mod;
  cout << (ans % mod + mod) % mod << "\n";
  return 0;
}