cp_library/alg/dp/butterfly/butterfly_masks_fn.py

• View this file on GitHub
• Last update: 2025-03-30 20:17:47+09:00

Verified with

• test/library-checker/set-power-series/subset_convolution_snippet.test.py

Code

import cp_library.__header__
import cp_library.alg.__header__
import cp_library.alg.dp.__header__
import cp_library.alg.dp.butterfly.__header__

def butterfly_masks(N, Z):
    for i in range(N):
        m = b = 1<<i
        while m < Z:
            yield m^b, m
            m = (m+1)|b

def fwht(A: list, N: int):
    for m0, m1 in butterfly_masks(N, len(A)):
        a0, a1 = A[m0], A[m1]
        A[m0], A[m1] = a0+a1, a0-a1
    return A

def subset_zeta(A: list[int], N: int):
    for m0, m1 in butterfly_masks(N, len(A)):
        A[m1] += A[m0]
    return A

def subset_zeta_pair(A: list[int], B: list[int], N: int):
    for m0, m1 in butterfly_masks(N, len(A)):
        A[m1] += A[m0]
        B[m1] += B[m0]
    return A, B

def subset_mobius(A: list[int], N: int):
    for m0, m1 in butterfly_masks(N, len(A)):
        A[m1] -= A[m0]
    return A

def superset_zeta(A, N: int):
    for m0, m1 in butterfly_masks(N, len(A)):
        A[m0] += A[m1]
    return A

def superset_mobius(A, N: int):
    for m0, m1 in butterfly_masks(N, len(A)):
        A[m0] -= A[m1]
    return A

def popcnts(N):
    P = [0]*(1 << N)
    for i in range(N):
        for m in range(b := 1<<i):
            P[m^b] = P[m] + 1
    return P


def subset_conv(A,B,N):
    assert len(A) == len(B)
    Z = (N+1)*(M := 1<<N)
    Ar,Br,Cr,P = [0]*Z, [0]*Z, [0]*Z, popcnts(N)
    for i,p in enumerate(P): Ar[p<<N|i], Br[p<<N|i] = A[i], B[i]
    subset_zeta_pair(Ar, Br, N)
    for i in range(0,Z,M):
        for j in range(0,Z-i,M):
            ij = i+j
            for k in range(M): Cr[ij|k] += Ar[i|k] * Br[j|k]
    subset_mobius(Cr, N)
    for i,p in enumerate(P): A[i] = Cr[p<<N|i]
    return A

'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
             https://kobejean.github.io/cp-library               
'''


'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
    x₀ ────────●─●────────●───●────────●───────●────────► X₀
                ╳          ╲ ╱          ╲     ╱          
    x₄ ────────●─●────────●─╳─●────────●─╲───╱─●────────► X₁
                           ╳ ╳          ╲ ╲ ╱ ╱          
    x₂ ────────●─●────────●─╳─●────────●─╲─╳─╱─●────────► X₂
                ╳          ╱ ╲          ╲ ╳ ╳ ╱          
    x₆ ────────●─●────────●───●────────●─╳─╳─╳─●────────► X₃
                                        ╳ ╳ ╳ ╳         
    x₁ ────────●─●────────●───●────────●─╳─╳─╳─●────────► X₄
                ╳          ╲ ╱          ╱ ╳ ╳ ╲          
    x₅ ────────●─●────────●─╳─●────────●─╱─╳─╲─●────────► X₅
                           ╳ ╳          ╱ ╱ ╲ ╲          
    x₃ ────────●─●────────●─╳─●────────●─╱───╲─●────────► X₆
                ╳          ╱ ╲          ╱     ╲          
    x₇ ────────●─●────────●───●────────●───────●────────► X₇
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
                 Algorithms - DP - Butterfly                     
'''

def butterfly_masks(N, Z):
    for i in range(N):
        m = b = 1<<i
        while m < Z:
            yield m^b, m
            m = (m+1)|b

def fwht(A: list, N: int):
    for m0, m1 in butterfly_masks(N, len(A)):
        a0, a1 = A[m0], A[m1]
        A[m0], A[m1] = a0+a1, a0-a1
    return A

def subset_zeta(A: list[int], N: int):
    for m0, m1 in butterfly_masks(N, len(A)):
        A[m1] += A[m0]
    return A

def subset_zeta_pair(A: list[int], B: list[int], N: int):
    for m0, m1 in butterfly_masks(N, len(A)):
        A[m1] += A[m0]
        B[m1] += B[m0]
    return A, B

def subset_mobius(A: list[int], N: int):
    for m0, m1 in butterfly_masks(N, len(A)):
        A[m1] -= A[m0]
    return A

def superset_zeta(A, N: int):
    for m0, m1 in butterfly_masks(N, len(A)):
        A[m0] += A[m1]
    return A

def superset_mobius(A, N: int):
    for m0, m1 in butterfly_masks(N, len(A)):
        A[m0] -= A[m1]
    return A

def popcnts(N):
    P = [0]*(1 << N)
    for i in range(N):
        for m in range(b := 1<<i):
            P[m^b] = P[m] + 1
    return P


def subset_conv(A,B,N):
    assert len(A) == len(B)
    Z = (N+1)*(M := 1<<N)
    Ar,Br,Cr,P = [0]*Z, [0]*Z, [0]*Z, popcnts(N)
    for i,p in enumerate(P): Ar[p<<N|i], Br[p<<N|i] = A[i], B[i]
    subset_zeta_pair(Ar, Br, N)
    for i in range(0,Z,M):
        for j in range(0,Z-i,M):
            ij = i+j
            for k in range(M): Cr[ij|k] += Ar[i|k] * Br[j|k]
    subset_mobius(Cr, N)
    for i,p in enumerate(P): A[i] = Cr[p<<N|i]
    return A

Back to top page