XOR・AND・OR 畳み込み

xor 畳み込みはアダマール変換、and 畳み込みと or 畳み込みは上位集合・部分集合についてのゼータ変換を使い、要素ごとの積に直してから逆変換で戻します。

• xor_convolution(f, g, mod): $h(k) = \sum_{i \oplus j = k} f(i) g(j)$ のリストを返す
• and_convolution(f, g, mod): $h(k) = \sum_{i \mathbin{\mathrm{and}} j = k} f(i) g(j)$ のリストを返す
• or_convolution(f, g, mod): $h(k) = \sum_{i \mathbin{\mathrm{or}} j = k} f(i) g(j)$ のリストを返す
• 計算量: どれも $O(N 2^N)$

f と g は同じ長さ（ $2$ のべき）にします。mod を渡すと、各値を mod で割った余りで返します（省略すると割らない）。xor 畳み込みで mod を渡すときは $2^N$ の逆元を使うので、mod は奇数にします。

部分集合や上位集合にわたる和だけが欲しいときは、高速ゼータ変換・メビウス変換 を使います。

コード

def xor_convolution(f, g, mod=0):
    # h[k] = (i xor j = k となるすべての組についての f[i] g[j] の和)。長さは同じ 2 べき
    n = len(f)
    f, g = list(f), list(g)
    for a in (f, g):
        b = 1
        while b < n:
            for i in range(0, n, 2 * b):
                for j in range(i, i + b):
                    a[j], a[j + b] = a[j] + a[j + b], a[j] - a[j + b]
            b *= 2
    h = [x * y % mod for x, y in zip(f, g)] if mod else [x * y for x, y in zip(f, g)]
    b = 1
    while b < n:
        for i in range(0, n, 2 * b):
            for j in range(i, i + b):
                h[j], h[j + b] = h[j] + h[j + b], h[j] - h[j + b]
        b *= 2
    if mod:
        inv = pow(n, -1, mod)
        return [x * inv % mod for x in h]
    return [x // n for x in h]


def and_convolution(f, g, mod=0):
    # h[k] = (i and j = k となるすべての組についての f[i] g[j] の和)。長さは同じ 2 べき
    n = len(f)
    f, g = list(f), list(g)
    for a in (f, g):
        b = 1
        while b < n:
            for i in range(0, n, 2 * b):
                for j in range(i, i + b):
                    a[j] += a[j + b]
            b *= 2
    h = [x * y % mod for x, y in zip(f, g)] if mod else [x * y for x, y in zip(f, g)]
    b = 1
    while b < n:
        for i in range(0, n, 2 * b):
            for j in range(i, i + b):
                h[j] -= h[j + b]
        b *= 2
    return [x % mod for x in h] if mod else h


def or_convolution(f, g, mod=0):
    # h[k] = (i or j = k となるすべての組についての f[i] g[j] の和)。長さは同じ 2 べき
    n = len(f)
    f, g = list(f), list(g)
    for a in (f, g):
        b = 1
        while b < n:
            for i in range(0, n, 2 * b):
                for j in range(i, i + b):
                    a[j + b] += a[j]
            b *= 2
    h = [x * y % mod for x, y in zip(f, g)] if mod else [x * y for x, y in zip(f, g)]
    b = 1
    while b < n:
        for i in range(0, n, 2 * b):
            for j in range(i, i + b):
                h[j + b] -= h[j]
        b *= 2
    return [x % mod for x in h] if mod else h