P5591 小猪佩奇学数学

P5591 小猪佩奇学数学

单位根反演：
$$
[k | j] = \frac {\sum _ {d = 0} ^ {k - 1} \omega _ k ^ {dj}} k
$$
那么
$$
\begin {aligned}
\left \lfloor \frac i k \right \rfloor & = \sum _ {j = 1} ^ i [k | j] \\
& = (\sum _ {j = 0} ^ i [k | j]) - 1 \\
& = (\sum _ {j = 0} ^ i \frac {\sum _ {d = 0} ^ {k - 1} \omega _ k ^ {dj}} k) - 1 \\
& = \frac {\sum _ {d = 0} ^ {k - 1} \sum _ {j = 0} ^ i \omega _ k ^ {dj}} k - 1 \\
\end {aligned}
$$

根据二项式定理：
$$
\sum _ {i = 0} ^ n \binom n i p ^ i = (p + 1) ^ n
$$

考虑如何使得单位根的幂都和组合数转化为右边的形式。

$$
\begin {aligned}
ANS = & \sum _ {i = 0} ^ n \binom n i p ^ i \left \lfloor \frac i k \right \rfloor \\
= & \sum _ {i = 0} ^ n \binom n i p ^ i (\frac {\sum _ {d = 0} ^ {k - 1} \sum _ {j = 0} ^ i \omega _ k ^ {dj}} k - 1) \\
= & \frac {\sum _ {d = 0} ^ {k - 1} \sum _ {i = 0} ^ n \binom n i p ^ i \sum _ {j = 0} ^ i \omega _ k ^ {dj}} k - \sum _ {i = 0} ^ n \binom n i p ^ i \\
= & \frac {\sum _ {d = 0} ^ {k - 1} \sum _ {i = 0} ^ n \binom n i p ^ i \frac {\omega _ k ^ {d (i + 1)} - 1} {\omega _ k ^ d - 1}} k - (p + 1) ^ n \\
= & \frac {\sum _ {d = 0} ^ {k - 1} \frac { \sum _ {i = 0} ^ n \binom n i p ^ i \omega _ k ^ {d (i + 1)} - \sum _ {i = 0} ^ n \binom n i p ^ i} {\omega _ k ^ d - 1}} k - (p + 1) ^ n \\
= & \frac {\sum _ {d = 0} ^ {k - 1} \frac {\omega _ k ^ d \sum _ {i = 0} ^ n \binom n i {(p \omega _ k ^ d)} ^ i - \sum _ {i = 0} ^ n \binom n i p ^ i} {\omega _ k ^ d - 1}} k - (p + 1) ^ n \\
= & \frac {\sum _ {d = 0} ^ {k - 1} \frac {\omega _ k ^ d (p \omega _ k ^ d + 1) ^ n - (p + 1) ^ n} {\omega _ k ^ d - 1}} k - (p + 1) ^ n \\
\end {aligned}
$$

注意到 $d = 0$ 时，计算等比数列时候公式应用时错误的，导致 $w _ k ^ d - 1 = 0$ 分母为 $0$ ，所以将其单独计算：
$$
\begin {aligned}
& \frac {\sum _ {i = 0} ^ n \binom n i p ^ i \sum _ {j = 0} ^ i \omega _ k ^ {0}} k \\
= &\frac { \sum _ {i = 0} ^ n \binom n i p ^ i (i + 1)} k \\
= &\frac { \sum _ {i = 1} ^ n \binom n i p ^ i i + (p + 1) ^ n} k \\
= &\frac { \sum _ {i = 1} ^ n n \binom {n - 1} {i - 1} p ^ i+ (p + 1) ^ n} k \\
= &\frac {np \sum _ {i = 1} ^ n\binom {n - 1} {i - 1} p ^ {i - 1}+ (p + 1) ^ n} k \\
= &\frac {np \sum _ {i = 0} ^ {n - 1}\binom {n - 1} i p ^ i + (p + 1) ^ n} k \\
= &\frac {np (p + 1) ^ {n - 1} + (p + 1) ^ n} k \\
\end {aligned}
$$

所以最终答案：
$$
\frac {np (p + 1) ^ {n - 1} + (p + 1) ^ n + \sum _ {d = 1} ^ {k - 1} \frac {\omega _ k ^ d (p \omega _ k ^ d + 1) ^ n - (p + 1) ^ n} {\omega _ k ^ d - 1}} k - (p + 1) ^ n \\
$$
因为 $k$ 很小，所以复杂度很优。

查看代码

#include <cstdio>
using namespace std;
template <class Type>
void read(Type &x)
{
    char c;
    bool flag = false;
    while ((c = getchar()) < '0' || c > '9')
        c == '-' && (flag = true);
    x = c - '0';
    while ((c = getchar()) >= '0' && c <= '9')
        x = (x << 3) + (x << 1) + c - '0';
    if (flag) x = ~x + 1;
}
template <class Type, class ...rest>
void read(Type &x, rest &...y) { read(x), read(y...); }
template <class Type>
void write(Type x)
{
    if (x < 0) putchar('-'), x = ~x + 1;
    if (x > 9) write(x / 10);
    putchar(x % 10 + '0');
}
typedef long long LL;
const int N = 30, mod = 998244353;
void adj (int &x) { x += x >> 31 & mod; }
int binpow (int b, int k = mod - 2)
{
    int res = 1;
    for (; k; k >>= 1, b = (LL)b * b % mod)
        if (k & 1) res = (LL)res * b % mod;
    return res;
}
struct ModInt
{
    int x;
    ModInt (int _ = 0) { adj(x = _); }
    int operator () () const { return x; }
    ModInt& operator += (const ModInt &_) { adj(x += _.x - mod); return *this; }
    ModInt& operator -= (const ModInt &_) { adj(x -= _.x); return *this; }
    ModInt& operator *= (const ModInt &_) { x = (LL)x * _.x % mod; return *this; }
    ModInt& operator /= (const ModInt &_) { x = (LL)x * binpow(_.x) % mod; return *this; }
    ModInt& operator ^= (const ModInt &_) { x = binpow(x, _.x); return *this; }
    ModInt operator + (const ModInt &_) const { ModInt res = x; res += _; return res; }
    ModInt operator - (const ModInt &_) const { ModInt res = x; res -= _; return res; }
    ModInt operator * (const ModInt &_) const { ModInt res = x; res *= _; return res; }
    ModInt operator / (const ModInt &_) const { ModInt res = x; res /= _; return res; }
    ModInt operator ^ (const ModInt &_) const { ModInt res = x; res ^= _; return res; }
};
int n, p, k;
int main ()
{
    read(n, p, k);
    ModInt res = (((ModInt)p + 1) ^ (n - 1)) * n * p + (((ModInt)p + 1) ^ n), t = binpow(3, (mod - 1) / k), s = t;
    for (int d = 1; d < k; ++d, s *= t)
        res += ((((ModInt)s * p + 1) ^ n) * s - (((ModInt)p + 1) ^ n)) / (s - 1);
    res /= k;
    res -= ((ModInt)p + 1) ^ n;
    write(res());
    return 0;
}