約数系ゼータ変換・メビウス変換

素数 $p$ ごとに「 $f(i)$ を $f(ip)$ に足す」ことを繰り返すと、約数にわたる和になります。高速ゼータ変換 でビットを 1 つずつ見るところを、素因数の指数を 1 つずつ見るのに変えたものです。リストの添字は $1$ から $N$ までを使い、 $0$ 番目は使いません。

• divisor_zeta(f, mod): $g(n) = \sum_{d \mid n} f(d)$ のリストを返す
• divisor_mobius(g, mod): divisor_zeta の逆変換
• multiple_zeta(f, mod): $g(n) = \sum_{n \mid m} f(m)$（ $m \le N$）のリストを返す
• multiple_mobius(g, mod): multiple_zeta の逆変換
• 計算量: どれも $O(N \log \log N)$

mod を渡すと、各値を mod で割った余りで返します（省略すると割らない）。渡したリストは書き換えません。 $f(1) = 1$ でそれ以外が $0$ のリストに divisor_mobius を使うと、メビウス関数 $\mu(n)$ の表になります。

例（gcd 畳み込み $h(k) = \sum_{\gcd(i, j) = k} f(i) g(j)$）:

F, G = multiple_zeta(f), multiple_zeta(g)
h = multiple_mobius([x * y for x, y in zip(F, G)])

lcm 畳み込み（ $\mathrm{lcm}(i, j) = k \le N$ となる組について足す）は、divisor_zeta と divisor_mobius で同じようにします。

コード

def _primes(n):
    # n 以下の素数のリスト
    s = [True] * (n + 1)
    res = []
    for i in range(2, n + 1):
        if s[i]:
            res.append(i)
            for j in range(i * i, n + 1, i):
                s[j] = False
    return res


def divisor_zeta(f, mod=0):
    # g[n] = (n の約数 d すべてについての f[d] の和)。添字は 1 から N（f[0] は使わない）
    f = [x % mod for x in f] if mod else list(f)
    n = len(f) - 1
    for p in _primes(n):
        for i in range(1, n // p + 1):
            f[i * p] += f[i]
    return [x % mod for x in f] if mod else f


def divisor_mobius(g, mod=0):
    # divisor_zeta の逆変換
    g = [x % mod for x in g] if mod else list(g)
    n = len(g) - 1
    for p in _primes(n):
        for i in range(n // p, 0, -1):
            g[i * p] -= g[i]
    return [x % mod for x in g] if mod else g


def multiple_zeta(f, mod=0):
    # g[n] = (N 以下の n の倍数 m すべてについての f[m] の和)。添字は 1 から N（f[0] は使わない）
    f = [x % mod for x in f] if mod else list(f)
    n = len(f) - 1
    for p in _primes(n):
        for i in range(n // p, 0, -1):
            f[i] += f[i * p]
    return [x % mod for x in f] if mod else f


def multiple_mobius(g, mod=0):
    # multiple_zeta の逆変換
    g = [x % mod for x in g] if mod else list(g)
    n = len(g) - 1
    for p in _primes(n):
        for i in range(1, n // p + 1):
            g[i] -= g[i * p]
    return [x % mod for x in g] if mod else g