P5591 小猪佩奇学数学（单位根反演）

P5591 小猪佩奇学数学

$$\begin{gathered} \sum _{i=0}^n{(}_i^n)\times p^i\times \lfloor \frac{i}{k}\rfloor \\ \lfloor \frac{i}{k}\rfloor =\frac{i-i\%k}{k} \\ \frac{1}{k}\sum _{i=0}^n{(}_i^n)\times p^i\times (i-i\%k) \\ \frac{1}{k}\sum _{i=0}^n{(}_i^n)\times p^i\times i-\frac{1}{k}\sum _{i=0}^n{(}_i^n)\times p^i(i\mathrm{m}\mathrm{o}\mathrm{d}k) \end{gathered}$$
考虑前项化简：
$$\begin{gathered} \frac{1}{k}\sum _{i=0}^n{(}_i^n)\times p^i\times i \\ \sum _{i=0}^n{(}_i^n)\times i=\sum _{i=0}^n\frac{n!}{i!(n-i)!}i=\sum _{i=1}^n\frac{n!}{(i-1)!(n-i)!} \\ n\sum _{i=1}^n\frac{(n-1)!}{(i-1)!(n-i)!}=n\sum _{i=1}^n{(}_{i-1}^{n-1}) \\ 原式\frac{np}{k}\sum _{i=1}^n{(}_{i-1}^{n-1})p^{i-1}=\frac{np}{k}(p+1{)}^{n-1} \end{gathered}$$
考虑后项化简：
$$\begin{gathered} \frac{1}{k}\sum _{d=0}^{k-1}d\sum _{i=0}^n{(}_i^n)p^i[k\mid i-d] \\ 对[k\mid i-d]进行单位根反演 \\ \frac{1}{k}\sum _{d=0}^{k-1}d\sum _{i=0}^n{(}_i^n)\times p^i\times \frac{1}{k}\sum _{j=0}^{k-1}{w}_k^{j(i-d)} \\ \frac{1}{k^2}\sum _{d=0}^{k-1}d\sum _{j=0}^{k-1}{w}_k^{-jd}\sum _{i=0}^n{(}_i^n)p^i{w}_k^{ij} \\ \frac{1}{k^2}\sum _{d=0}^{k-1}d\sum _{j=0}^{k-1}{w}_k^{-jd}(p{w}_k^j+1{)}^n \\ \frac{1}{k^2}\sum _{j=0}^{k-1}(p\times w_k^j+1{)}^n\sum _{d=0}^{k-1}d\times w_k^{-jd} \end{gathered}$$
$\sum _{d=0}^{k-1}d\times w_k^{-jd}$形如$\sum _{i=0}^{n-1}i{r}^i$，考虑形如$\sum _{i=0}^{n-1}(i+1)r^{i+1}$这样的式子化简：
$$\begin{gathered} \sum _{i=0}^{n-1}(i+1)r^{i+1}=r\sum _{i=0}^{n-1}i{r}^i+\sum _{i=0}^{n-1}{r}^{i+1}=r\sum _{i=0}^{n-1}i{r}^i+\frac{r(r^n-1)}{r-1} \\ 有\sum _{i=0}^{n-1}i{r}^i=\sum _{i=0}^{n-1}(i+1)r^{i+1}-0r^0-n{r}^n \\ 所以\sum _{i=0}^{n-1}i{r}^i=r\sum _{i=0}^{n-1}i{r}^i+\frac{r(r^n-1)}{r-1}-n{r}^n \\ 整理可得\sum _{i=0}^{n-1}i{r}^i=\frac{n{r}^n-\frac{r(r^n-1)}{r-1}}{r-1}=\frac{n{r}^n(r-1)-r(r^n-1)}{(r-1{)}^2}=\frac{n{r}^n}{r-1}-\frac{r(r^n-1)}{(r-1{)}^2} \end{gathered}$$
特殊情况$r=1$时$\sum _{i=0}^{n-1}i{r}^i=\frac{n(n-1)}{2}$，对后项再进行化简：
$$\begin{gathered} \sum _{d=0}^{k-1}d\times w_k^{-jd} \\ 考虑j\ne 0时，也就是w_k^j不为1 \\ \sum _{d=0}^{k-1}d\times w_k^{-jd}=\frac{k{w}_k^{-jk}}{w_k^{-j}-1}-\frac{w_k^{-j}(w_k^{-jk}-1)}{(w_k^{-j}-1{)}^2}=\frac{k}{w_k^{-j}-1} \\ 后项整体有:\frac{1}{k^2}\sum _{j=0}^{k-1}(p\times w_k^j+1{)}^n\frac{k}{w_k^{-j}-1} \end{gathered}$$
最后答案为：
$$\begin{gathered} \frac{np}{k}(p+1{)}^{n-1}-\left(\frac{k(k-1)}{2}\frac{(p+1{)}^n}{k^2}+\sum _{j=1}^{k-1}\frac{(p\times w_k^j+1{)}^n}{k(w_k^{-j}-1)}\right) \\ \frac{np}{k}(p+1{)}^{n-1}-\left(\frac{(k-1)(p+1{)}^n}{2k}+\sum _{j=1}^{k-1}\frac{(p\times w_k^j+1{)}^n}{k(w_k^{-j}-1)}\right) \end{gathered}$$

#include <bits/stdc++.h>

using namespace std;

const int mod = 998244353, N = 2e6 + 10;

int w[N], n, p, k;

int quick_pow(int a, int n) {
  int ans = 1;
  while (n) {
    if (n & 1) {
      ans = 1ll * ans * a % mod;
    }
    a = 1ll * a * a % mod;
    n >>= 1;
  }
  return ans;
}

int inv(int x) {
  return quick_pow(x, mod - 2);
}

int main() {
  // freopen("in.txt", "r", stdin);
  // freopen("out.txt", "w", stdout);
  scanf("%d %d %d", &n, &p, &k);
  int wn = quick_pow(3, (mod - 1) / k);
  int ans = 1ll * n * p % mod * quick_pow(p + 1, n - 1) % mod * inv(k) % mod;
  int res = 1ll * (k - 1) * quick_pow(p + 1, n) % mod * inv(2 * k) % mod;
  w[0] = 1;
  for (int i = 1; i < k; i++) {
    w[i] = 1ll * w[i - 1] * wn % mod;
  }
  for (int i = 1; i < k; i++) {
    res = (res + 1ll * quick_pow(1ll * p * w[i] % mod + 1, n) * inv(1ll * k * (w[k - i] - 1) % mod) % mod) % mod;
  }
  ans = (ans - res + mod) % mod;
  printf("%d\n", ans);
  return 0;
}