cp_library/math/sps/mod/sps_composite_fn.py

View this file on GitHub
Last update: 2025-08-24 20:27:05+09:00

Depends on

Verified with

test/library-checker/set-power-series/polynomial_composite_set_power_series.test.py

Code

import cp_library.__header__
from cp_library.bit.popcnts_fn import popcnts
from cp_library.ds.list.elist_fn import elist
from cp_library.ds.view.view_cls import view
import cp_library.math.__header__
from cp_library.math.conv.ior_zeta_fn import ior_zeta
from cp_library.math.conv.ior_mobius_ranked_fn import ior_mobius_ranked
import cp_library.math.sps.__header__
import cp_library.math.sps.mod.__header__

def isubset_conv_zeta_ranked(Ar: list[int], Br: list[int], n: int, N: int, mod: int) -> list[int]:
    m = 1<<n
    for ij in range(n,-1,-1):
        ij_, i_ = (ij+1)<<N|m, ij<<n
        for k in range(m): Ar[ij_|k] = Br[i_|k] * Ar[k] % mod
        for i in range(ij):
            j = ij-i; i_, j_ = i<<n, j<<N
            for k in range(m): Ar[ij_|k] = (Ar[ij_|k] + Br[i_|k] * Ar[j_|k]) % mod

def sps_composite(A: list[int], B: list[int], mod: int) -> list[int]:
    C = [0]*(M := 1 << (N := len(B).bit_length() - 1))
    if not A: return C
    dA, B0, B1, Br, Cr, pcnt = A[:], elist(N+1), view(B), elist(N), [0]*(Z := (N+1)*M), popcnts(N)
    for n in range(N+1):
        if n < N:
            # zeta transform of ranked 
            B1.set_range(1<<n, 2<<n)
            br = [0]*(z := (n+1)*(m := 1<<n))
            for i in range(m): br[pcnt[i]<<n|i] = B1[i]
            ior_zeta(br, n)
            for i in range(z): br[i] %= mod
            Br.append(br)
        # evaluate current polynomial at B[0] using Horner's method
        t = 0
        for j in range(len(dA)-1, -1, -1): t = (t * B[0] + dA[j]) % mod
        B0.append(t)
        # update dA to be the derivative
        for j in range(1, len(dA)): dA[j-1] = (j * dA[j]) % mod
        if dA: dA[-1] = 0
    for n in range(N+1):
        for m in range(n-1, -1, -1):
            # effectively computes `C[1<<m:2<<m] = subset_conv(C[:1<<m], B[1<<m:2<<m])`
            # but basically maintains `Cr`, the ranked zeta transformed `C`
            # partial zeta updates need to be made after loop ends to propagate contributions
            isubset_conv_zeta_ranked(Cr, Br[m], m, N, mod)
        # partial zeta updates
        for m in range(n):
            b = 1 << m
            for j in range(m+1):
                j <<= N
                for k in range(j, j|b): Cr[k|b] += Cr[k]
        for k in range(1<<n): Cr[k] = B0[~n]
    ior_mobius_ranked(Cr, N, M, Z)
    for i, p in enumerate(pcnt): C[i] = Cr[p<<N|i] % mod
    return C

'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
             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



def elist(hint: int) -> list: ...
try:
    from __pypy__ import newlist_hint
except:
    def newlist_hint(hint): return []
elist = newlist_hint
    

from typing import Generic
from typing import TypeVar

_S = TypeVar('S'); _T = TypeVar('T'); _U = TypeVar('U'); _T1 = TypeVar('T1'); _T2 = TypeVar('T2'); _T3 = TypeVar('T3'); _T4 = TypeVar('T4'); _T5 = TypeVar('T5'); _T6 = TypeVar('T6')
import sys

def list_find(lst: list, value, start = 0, stop = sys.maxsize):
    try:
        return lst.index(value, start, stop)
    except:
        return -1


class view(Generic[_T]):
    __slots__ = 'A', 'l', 'r'
    def __init__(V, A: list[_T], l: int = 0, r: int = 0): V.A, V.l, V.r = A, l, r
    def __len__(V): return V.r - V.l
    def __getitem__(V, i: int): 
        if 0 <= i < V.r - V.l: return V.A[V.l+i]
        else: raise IndexError
    def __setitem__(V, i: int, v: _T): V.A[V.l+i] = v
    def __contains__(V, v: _T): return list_find(V.A, v, V.l, V.r) != -1
    def set_range(V, l: int, r: int): V.l, V.r = l, r
    def index(V, v: _T): return V.A.index(v, V.l, V.r) - V.l
    def reverse(V):
        l, r = V.l, V.r-1
        while l < r: V.A[l], V.A[r] = V.A[r], V.A[l]; l += 1; r -= 1
    def sort(V, /, *args, **kwargs):
        A = V.A[V.l:V.r]; A.sort(*args, **kwargs)
        for i,a in enumerate(A,V.l): V.A[i] = a
    def pop(V): V.r -= 1; return V.A[V.r]
    def append(V, v: _T): V.A[V.r] = v; V.r += 1
    def popleft(V): V.l += 1; return V.A[V.l-1]
    def appendleft(V, v: _T): V.l -= 1; V.A[V.l] = v; 
    def validate(V): return 0 <= V.l <= V.r <= len(V.A)

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

def ior_zeta(A: list[int], N: int, Z: int = None):
    Z = Z if Z else len(A)
    for i in range(N):
        m = b = 1<<i
        while m < Z: A[m] += A[m^b]; m = m+1|b
    return A

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_zeta_ranked(Ar: list[int], Br: list[int], n: int, N: int, mod: int) -> list[int]:
    m = 1<<n
    for ij in range(n,-1,-1):
        ij_, i_ = (ij+1)<<N|m, ij<<n
        for k in range(m): Ar[ij_|k] = Br[i_|k] * Ar[k] % mod
        for i in range(ij):
            j = ij-i; i_, j_ = i<<n, j<<N
            for k in range(m): Ar[ij_|k] = (Ar[ij_|k] + Br[i_|k] * Ar[j_|k]) % mod

def sps_composite(A: list[int], B: list[int], mod: int) -> list[int]:
    C = [0]*(M := 1 << (N := len(B).bit_length() - 1))
    if not A: return C
    dA, B0, B1, Br, Cr, pcnt = A[:], elist(N+1), view(B), elist(N), [0]*(Z := (N+1)*M), popcnts(N)
    for n in range(N+1):
        if n < N:
            # zeta transform of ranked 
            B1.set_range(1<<n, 2<<n)
            br = [0]*(z := (n+1)*(m := 1<<n))
            for i in range(m): br[pcnt[i]<<n|i] = B1[i]
            ior_zeta(br, n)
            for i in range(z): br[i] %= mod
            Br.append(br)
        # evaluate current polynomial at B[0] using Horner's method
        t = 0
        for j in range(len(dA)-1, -1, -1): t = (t * B[0] + dA[j]) % mod
        B0.append(t)
        # update dA to be the derivative
        for j in range(1, len(dA)): dA[j-1] = (j * dA[j]) % mod
        if dA: dA[-1] = 0
    for n in range(N+1):
        for m in range(n-1, -1, -1):
            # effectively computes `C[1<<m:2<<m] = subset_conv(C[:1<<m], B[1<<m:2<<m])`
            # but basically maintains `Cr`, the ranked zeta transformed `C`
            # partial zeta updates need to be made after loop ends to propagate contributions
            isubset_conv_zeta_ranked(Cr, Br[m], m, N, mod)
        # partial zeta updates
        for m in range(n):
            b = 1 << m
            for j in range(m+1):
                j <<= N
                for k in range(j, j|b): Cr[k|b] += Cr[k]
        for k in range(1<<n): Cr[k] = B0[~n]
    ior_mobius_ranked(Cr, N, M, Z)
    for i, p in enumerate(pcnt): C[i] = Cr[p<<N|i] % mod
    return C