Project Euler 840

发表于 分类于 阅读次数： 本文字数： 3.1k 阅读时长 ≈ 3 分钟

Project Euler 840

题目

Sum of Products

A partition of \(n\) is a set of positive integers for which the sum equals \(n\).

The partitions of \(5\) are:

\(\{5\},\{1,4\},\{2,3\},\{1,1,3\},\{1,2,2\},\{1,1,1,2\}\) and \(\{1,1,1,1,1\}\).

Further we define the function \(D(p)\) as:

\[\begin{align} \begin{split} D(1) &= 1 \\ D(p) &= 1, \text{ for any prime } p \\ D(pq) &= D(p)q + pD(q), \text{ for any positive integers } p,q \gt 1. \end{split} \end{align}\]

Now let \(\{a_1, a_2,\ldots,a_k\}\) be a partition of \(n\).

We assign to this particular partition the value:

\[P=\prod_{j=1}^{k}D(a_j). \]

\(G(n)\) is the sum of \(P\) for all partitions of \(n\).

We can verify that \(G(10) = 164\).

We also define:

\[S(N)=\sum_{n=1}^{N}G(n).\]

You are given \(S(10)=396\).

Find \(S(5\times 10^4) \mod 999676999\).

解决方案

令\(N=5\times 10^4,M=999676999\)。递推式\(D(pq)=D(p)q+pD(q)\)与微积分的乘积法则一致，配合 \(D(p)=1\)（素数）正是经典的算术导数（arithmetic derivative）的定义方式（除题目额外规定 \(D(1)=1\) 以外）。因此 \(D(n)\) 具备可由素因子分解快速计算的闭式。

对 \(e\ge 1\),\(D(p^e)=D(p^{e-1}\cdot p)=D(p^{e-1})p+p^{e-1}D(p)=pD(p^{e-1})+p^{e-1}.\)由此递推可得\(D(p^e)=ep^{e-1}.\)

对于一般情况，设\(n\)的质因子分解为\(\displaystyle{n=\prod_{i=1}^t p_i^{e_i}.}\)利用乘积法则反复展开（或用对数微分的等价形式），可得标准公式：\(\displaystyle{D(n)=n\sum_{i=1}^t \frac{e_i}{p_i}}\)等价地写成无分数的形式：

\[ D(n)=\sum_{i=1}^t e_i\cdot \frac{n}{p_i}. \]

因此只要知道 \(n\) 的素因子及指数，就能在 \(O(\omega(n))\) 时间算出 \(D(n)\)（\(\omega\) 为不同素因子个数）。

把分拆按每个部件大小 \(k\) 出现次数 \(m_k\) 表示，则\(\displaystyle{\sum_{k\ge 1} km_k=n,P=\prod_{k\ge 1} D(k)^{m_k}.}\)

对固定的 \(k\)，出现次数 \(m\) 可以是 \(0,1,2,\dots\)，对生成函数的贡献是几何级数：\(\displaystyle{\sum_{m\ge 0}\bigl(D(k)x^k\bigr)^m=\frac{1}{1-D(k)x^k}.}\)不同 \(k\) 相互独立相乘，于是得到总体生成函数\(\displaystyle{ F(x)=\sum_{n\ge 0}G(n)x^n=\prod_{k\ge 1}\frac{1}{1-D(k)x^k}.}\)

我们只需要到 \(N\) 的系数，因此在模 \(x^{N+1}\) 意义下可截断为

\[ F(x)\equiv \prod_{k=1}^{N}\frac{1}{1-D(k)x^k}\pmod{x^{N+1}}. \]

其中常数项 \(G(0)=1\) 对应空分拆。最终目标和为

\[ S(N)=\sum_{n=1}^N G(n)=\left(\sum_{n=0}^N [x^n]F(x)\right)-1. \]

可见，如果直接做完全背包型 DP：\(G(n)\leftarrow G(n)+D(k)G(n-k)\)是 \(O(N^2)\)的时间复杂度，比较慢。关键是对生成函数取对数：

\[ \log F(x)=\sum_{k=1}^N -\log(1-D(k)x^k) =\sum_{k=1}^N \sum_{m\ge 1}\frac{D(k)^m}{m}x^{km}. \]

令\(\displaystyle{L(x)=\log F(x)=\sum_{t=1}^N L_t x^t,}\)则系数满足\(\displaystyle{L_t=\sum_{k\mid t}\frac{D(k)^{t/k}}{t/k}.}\)（也就是令 \(t=km\)，对所有 \(k\) 枚举 \(m=t/k\)）

于是只要先算出 \(L(x)\) 的前 \(N\) 项，就有\(F(x)=\exp(L(x)).\)所以核心变成：在模 \(M\) 下做形式幂级数的 \(\exp\)。

使用模逆元，我们可以在 \(O(N\log N)\)的时间计算出\(L(x)\bmod x^{N+1}\)的表达式。

接下来我们假设下面所有多项式运算都在模 \(x^m\) 截断下进行。

对于多项式求逆，即给定 \(f(0)\ne 0\)，求 \(g=f^{-1}\pmod{x^m}\)。通过 Newton 迭代：\(g_i=g_{i-1}\cdot (2-fg_{i-1})\pmod{x^{2n}}.\)从 \(g_1=f(0)^{-1}\) 开始，每次精度翻倍，直到达到 \(m\)。

对于多项式对数，要求 \(f(0)=1\)，并定义\(\displaystyle{\log f=\int \frac{f'}{f}dx.}\)实现步骤如下：求导 \(f'\)；求逆 \(f^{-1}\)；卷积得到 \(f'\cdot f^{-1}\)；并积分即可。

对于多项式指数，给定幂级数 \(a(x)\) 且满足 \(a_0=0\)，要求 \(f=\exp(a)\)（此时必有 \(f_0=1\)）。Newton 迭代的标准形式为：\(f_{i}=f_{i-1}\cdot \bigl(1+a-\log f_{i-1}\bigr)\pmod{x^{2n}}.\)从 \(f=1\) 开始精度翻倍到 \(N+1\)。在本题中，取 \(a(x)=L(x)=\log F(x)\)，因此最终 \(F(x)=\exp(L(x))\)。

这些步骤的瓶颈全部是多项式乘法（卷积）。

接下来使用 FFT 做模数乘法。可见，模数 \(M\) 不是常见 NTT 友好模数时，可以用实数 FFT 做卷积，再通过“拆位”避免精度溢出。把系数表示为：\(a = a_0 + a_1 B, b=b_0+b_1 B,\)其中 \(B=2^{15}\)，\(a_0,a_1,b_0,b_1\) 都在较小范围内。

则\(ab = (a_0b_0) + (a_0b_1+a_1b_0)B + (a_1b_1)B^2.\)分别用 FFT 做三次实数卷积，四舍五入取整后再按模数重建即可。再配合一个小规模阈值在规模很小时改用朴素乘法，整体速度更稳。

最终，计算 \(F(x)=\exp(L(x))\pmod{x^{N+1}}\)后，得到 \(G(n)=[x^n]F(x)\)。

代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
import numpy as np

from tools import Factorizer

N = 50000
MOD = 999676999



def conv_mod_fft(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    na = a.size
    nb = b.size
    if na == 0 or nb == 0:
        return np.zeros(0, dtype=np.int64)
    need = na + nb - 1
    L = 1 << ((need - 1).bit_length())
    base = 1 << 15

    a = a.astype(np.int64, copy=False) % MOD
    b = b.astype(np.int64, copy=False) % MOD

    a0 = (a % base).astype(np.float64, copy=False)
    a1 = (a // base).astype(np.float64, copy=False)
    b0 = (b % base).astype(np.float64, copy=False)
    b1 = (b // base).astype(np.float64, copy=False)

    fa0 = np.fft.rfft(a0, L)
    fa1 = np.fft.rfft(a1, L)
    fb0 = np.fft.rfft(b0, L)
    fb1 = np.fft.rfft(b1, L)

    c00 = np.fft.irfft(fa0 * fb0, L)[:need]
    c01 = np.fft.irfft(fa0 * fb1 + fa1 * fb0, L)[:need]
    c11 = np.fft.irfft(fa1 * fb1, L)[:need]

    c00 = np.rint(c00).astype(np.int64) % MOD
    c01 = np.rint(c01).astype(np.int64) % MOD
    c11 = np.rint(c11).astype(np.int64) % MOD

    bb = base % MOD
    bb2 = (base * base) % MOD
    return (c00 + c01 * bb + c11 * bb2) % MOD


def poly_mul(a: np.ndarray, b: np.ndarray, lim: int) -> np.ndarray:
    if lim <= 0:
        return np.zeros(0, dtype=np.int64)
    if a.size == 0 or b.size == 0:
        return np.zeros(lim, dtype=np.int64)
    if a.size * b.size <= 4000:
        L = min(lim, a.size + b.size - 1)
        r = np.zeros(L, dtype=np.int64)
        aa = a.astype(np.int64, copy=False) % MOD
        bb = b.astype(np.int64, copy=False) % MOD
        for i in range(aa.size):
            ai = int(aa[i])
            if ai == 0:
                continue
            jmax = min(bb.size, L - i)
            if jmax > 0:
                r[i:i + jmax] = (r[i:i + jmax] + ai * bb[:jmax]) % MOD
        if L < lim:
            r = np.pad(r, (0, lim - L))
        return r[:lim]
    c = conv_mod_fft(a, b)
    if c.size < lim:
        c = np.pad(c, (0, lim - c.size))
    return c[:lim]


def poly_inv(f: np.ndarray, m: int) -> np.ndarray:
    g = np.array([pow(int(f[0]), MOD - 2, MOD)], dtype=np.int64)
    n = 1
    while n < m:
        n2 = min(2 * n, m)
        fg = poly_mul(f[:n2], g, n2)
        t = np.zeros(n2, dtype=np.int64)
        t[0] = (2 - fg[0]) % MOD
        if n2 > 1:
            t[1:] = (-fg[1:]) % MOD
        g = poly_mul(g, t, n2)
        n = n2
    return g[:m]


def poly_log(f: np.ndarray, m: int, inv: np.ndarray) -> np.ndarray:
    df = np.zeros(max(0, m - 1), dtype=np.int64)
    up = min(f.size, m)
    if up > 1:
        idx = np.arange(1, up, dtype=np.int64)
        df[:up - 1] = (f[1:up] * idx) % MOD
    invf = poly_inv(f, m)
    q = poly_mul(df, invf, m - 1)
    res = np.zeros(m, dtype=np.int64)
    if m > 1:
        res[1:] = (q[:m - 1] * inv[1:m]) % MOD
    return res


def poly_exp(a: np.ndarray, m: int, inv: np.ndarray) -> np.ndarray:
    f = np.array([1], dtype=np.int64)
    n = 1
    while n < m:
        n2 = min(2 * n, m)
        ln_f = poly_log(f, n2, inv)
        diff = (a[:n2] - ln_f) % MOD
        diff[0] = (diff[0] + 1) % MOD
        f = poly_mul(f, diff, n2)
        n = n2
    return f[:m]

def solve(N):
    inv = np.zeros(N + 1, dtype=np.int64)
    inv[1] = 1
    for i in range(2, N + 1):
        inv[i] = (MOD - (MOD // i) * inv[MOD % i] % MOD) % MOD

    spf = Factorizer(N).spf
    D = [0] * (N + 1)
    D[1] = 1
    for x in range(2, N + 1):
        t = x
        res = 0
        while t > 1:
            p = spf[t]
            e = 0
            while t % p == 0:
                t //= p
                e += 1
            res += e * (x // p)
        D[x] = res % MOD

    L = np.zeros(N + 1, dtype=np.int64)
    for k in range(1, N + 1):
        ak = D[k] % MOD
        if ak == 0:
            continue
        pw = ak
        lim = N // k
        for m in range(1, lim + 1):
            L[k * m] = (L[k * m] + pw * inv[m]) % MOD
            pw = (pw * ak) % MOD

    f = poly_exp(L, N + 1, inv)
    ans = int(f[1:].sum() % MOD)
    return ans


if __name__ == "__main__":
    print(solve(N))