cp-library

This documentation is automatically generated by online-judge-tools/verification-helper

View the Project on GitHub kobejean/cp-library

:warning: cp_library/math/sps/mod/sps_mul_fn.py

Depends on

Code

import cp_library.__header__
import cp_library.math.__header__
from cp_library.math.conv.mod.subset_conv_fn import subset_conv
import cp_library.math.sps.__header__

def sps_mul(A, B, mod):
    N = len(A).bit_length() - 1
    return subset_conv(A, B, N, mod)
'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
             https://kobejean.github.io/cp-library               
'''



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
'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
    x₀ ────────●─●────────●───●────────●───────●────────► X₀
                ╳          ╲ ╱          ╲     ╱          
    x₄ ────────●─●────────●─╳─●────────●─╲───╱─●────────► X₁
                           ╳ ╳          ╲ ╲ ╱ ╱          
    x₂ ────────●─●────────●─╳─●────────●─╲─╳─╱─●────────► X₂
                ╳          ╱ ╲          ╲ ╳ ╳ ╱          
    x₆ ────────●─●────────●───●────────●─╳─╳─╳─●────────► X₃
                                        ╳ ╳ ╳ ╳         
    x₁ ────────●─●────────●───●────────●─╳─╳─╳─●────────► X₄
                ╳          ╲ ╱          ╱ ╳ ╳ ╲          
    x₅ ────────●─●────────●─╳─●────────●─╱─╳─╲─●────────► X₅
                           ╳ ╳          ╱ ╱ ╲ ╲          
    x₃ ────────●─●────────●─╳─●────────●─╱───╲─●────────► X₆
                ╳          ╱ ╲          ╱     ╲          
    x₇ ────────●─●────────●───●────────●───────●────────► X₇
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
                      Math - Convolution                     
'''



def max2(a, b): return a if a > b else b


def ior_zeta_pair_ranked(A, B, N, M, Z):
    for i in range(0, Z, M):
        l, r = i+(1<<(i>>N))-1, i+M
        for j in range(N):
            m = l|(b := 1<<j)
            while m < r: A[m] += A[m^b]; B[m] += B[m^b]; m = m+1|b
    return A, B

def ior_mobius_ranked(A: list[int], N: int, M: int, Z: int):
    for i in range(0, Z, M):
        l, r = i, i+M-(1<<(N-(i>>N)))+1
        for j in range(N):
            m = l|(b := 1<<j)
            while m < r: A[m] -= A[m^b]; m = m+1|b
    return A

def isubset_conv_ranked(Ar, Br, N, M, Z, mod) -> list[int]:
    ior_zeta_pair_ranked(Ar, Br, N, M, Z)
    for i in range(Z): Ar[i], Br[i] = Ar[i]%mod, Br[i]%mod
    for ij in range(Z-M,-1,-M):
        for k in range(M): Ar[ij|k] = (Ar[ij|k] * Br[k]) % mod
        r = M-(1 << (N-(ij>>N)))+1
        for i in range(0,ij,M):
            j = ij-i; l = (1 << (max2(i,j)>>N))-1
            for k in range(l,r): Ar[ij|k] += Ar[i|k] * Br[j|k] % mod
    return ior_mobius_ranked(Ar, N, M, Z)

def subset_conv(A: list[int], B: list[int], N: int, mod: int) -> list[int]:
    Z = (N+1)*(M:=1<<N)
    Ar, Br, C, P = [0]*Z, [0]*Z, [0]*M, popcnts(N)
    for i, p in enumerate(P): Ar[p<<N|i], Br[p<<N|i] = A[i], B[i]
    isubset_conv_ranked(Ar, Br, N, M, Z, mod)
    for i, p in enumerate(P): C[i] = Ar[p<<N|i] % mod
    return C


def sps_mul(A, B, mod):
    N = len(A).bit_length() - 1
    return subset_conv(A, B, N, mod)
Back to top page