高速ゼータ変換・メビウス変換

集合 $S \subseteq \lbrace 0, 1, \ldots, N - 1 \rbrace$ をビットで表した整数を添字として、長さ $2^N$ のリスト f で $f(S)$ を表します。1 ビットずつ順に足し込むことで、すべての $S$ について和をまとめて求めます。

mod を渡すと、各値を mod で割った余りで返します（省略すると割らない）。渡したリストは書き換えません。

例（OR 畳み込み $h(S) = \sum_{T \cup U = S} f(T) g(U)$）:

h = subset_mobius([x * y for x, y in zip(subset_zeta(f), subset_zeta(g))])

AND 畳み込み（ $T \cap U = S$ について足す）は、superset_zeta と superset_mobius で同じようにします。

コード

def subset_zeta(f, mod=0):
    # g[S] = (S の部分集合 T すべてについての f[T] の和)
    f = [x % mod for x in f] if mod else list(f)
    n = len(f)
    b = 1
    while b < n:
        for i in range(0, n, 2 * b):
            for j in range(i, i + b):
                f[j + b] += f[j]
        if mod:
            f = [x % mod for x in f]
        b *= 2
    return f


def subset_mobius(g, mod=0):
    # subset_zeta の逆変換
    g = [x % mod for x in g] if mod else list(g)
    n = len(g)
    b = 1
    while b < n:
        for i in range(0, n, 2 * b):
            for j in range(i, i + b):
                g[j + b] -= g[j]
        if mod:
            g = [x % mod for x in g]
        b *= 2
    return g


def superset_zeta(f, mod=0):
    # g[S] = (S を含む集合 T すべてについての f[T] の和)
    f = [x % mod for x in f] if mod else list(f)
    n = len(f)
    b = 1
    while b < n:
        for i in range(0, n, 2 * b):
            for j in range(i, i + b):
                f[j] += f[j + b]
        if mod:
            f = [x % mod for x in f]
        b *= 2
    return f


def superset_mobius(g, mod=0):
    # superset_zeta の逆変換
    g = [x % mod for x in g] if mod else list(g)
    n = len(g)
    b = 1
    while b < n:
        for i in range(0, n, 2 * b):
            for j in range(i, i + b):
                g[j] -= g[j + b]
        if mod:
            g = [x % mod for x in g]
        b *= 2
    return g