cp_library/math/table/primes_cls.py

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

Depends on

• cp_library/ds/reserve_fn.py

Required by

• cp_library/math/conv/gcd_conv_fn.py
• cp_library/math/conv/lcm_conv_fn.py
• cp_library/math/table/linear_sieve_cls.py
• cp_library/math/table/linear_sieve_cnts_cls.py
• cp_library/math/table/sieve_cls.py
• cp_library/math/table/sieve_proto.py
• cp_library/math/table/totient_cls.py

Verified with

• test/library-checker/convolution/gcd_convolution.test.py
• test/library-checker/convolution/lcm_convolution.test.py

Code

import cp_library.math.table.__header__
import operator
from typing import Callable
from cp_library.ds.reserve_fn import reserve

class Primes(list[int]):
    def __init__(P, N: int):
        super().__init__()
        spf = [0] * (N + 1)
        spf[0], spf[1] = 0, 1
        reserve(P, N)

        for i in range(2, N + 1):
            if spf[i] == 0:
                spf[i] = i
                P.append(i)
            for p in P:
                if p > spf[i] or i*p > N: break
                spf[i*p] = p
        P.spf = spf

    def divisor_zeta(P, A: list[int], op: Callable[[int,int], int] = operator.add) -> list[int]:
        N = len(A)-1
        for p in P:
            for i in range(1, N//p+1): A[i*p] = op(A[i*p], A[i])
        return A
    
    def divisor_mobius(P, A: list[int], diff: Callable[[int,int], int] = operator.sub) -> list[int]:
        N = len(A)-1
        for p in P:
            for i in range(N//p, 0, -1): A[i*p] = diff(A[i*p], A[i])
        return A
    
    def multiple_zeta(P, A: list[int], op: Callable[[int,int], int] = operator.add) -> list[int]:
        N = len(A)-1
        for p in P:
            for i in range(N//p, 0, -1): A[i] = op(A[i], A[i*p])
        return A
    
    def multiple_mobius(P, A: list[int], diff: Callable[[int,int], int] = operator.sub) -> list[int]:
        N = len(A)-1
        for p in P:
            for i in range(1, N//p+1): A[i] = diff(A[i], A[i*p])
        return A
    
    def gcd_conv(P, A: list[int], B: list[int], add = operator.add, sub = operator.sub, mul = operator.mul):
        A, B = P.multiple_zeta(A, add), P.multiple_zeta(B, add)
        for i, b in enumerate(B): A[i] = mul(A[i], b)
        return P.multiple_mobius(A, sub)
    
    def lcm_conv(P, A: list[int], B: list[int], add = operator.add, sub = operator.sub, mul = operator.mul):
        A, B = P.divisor_zeta(A, add), P.divisor_zeta(B, add)
        for i, b in enumerate(B): A[i] = mul(A[i], b)
        return P.divisor_mobius(A, sub)

'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
             https://kobejean.github.io/cp-library               
'''
import operator
from typing import Callable


def reserve(A: list, est_len: int) -> None: ...
try:
    from __pypy__ import resizelist_hint
except:
    def resizelist_hint(A: list, est_len: int):
        pass
reserve = resizelist_hint

class Primes(list[int]):
    def __init__(P, N: int):
        super().__init__()
        spf = [0] * (N + 1)
        spf[0], spf[1] = 0, 1
        reserve(P, N)

        for i in range(2, N + 1):
            if spf[i] == 0:
                spf[i] = i
                P.append(i)
            for p in P:
                if p > spf[i] or i*p > N: break
                spf[i*p] = p
        P.spf = spf

    def divisor_zeta(P, A: list[int], op: Callable[[int,int], int] = operator.add) -> list[int]:
        N = len(A)-1
        for p in P:
            for i in range(1, N//p+1): A[i*p] = op(A[i*p], A[i])
        return A
    
    def divisor_mobius(P, A: list[int], diff: Callable[[int,int], int] = operator.sub) -> list[int]:
        N = len(A)-1
        for p in P:
            for i in range(N//p, 0, -1): A[i*p] = diff(A[i*p], A[i])
        return A
    
    def multiple_zeta(P, A: list[int], op: Callable[[int,int], int] = operator.add) -> list[int]:
        N = len(A)-1
        for p in P:
            for i in range(N//p, 0, -1): A[i] = op(A[i], A[i*p])
        return A
    
    def multiple_mobius(P, A: list[int], diff: Callable[[int,int], int] = operator.sub) -> list[int]:
        N = len(A)-1
        for p in P:
            for i in range(1, N//p+1): A[i] = diff(A[i], A[i*p])
        return A
    
    def gcd_conv(P, A: list[int], B: list[int], add = operator.add, sub = operator.sub, mul = operator.mul):
        A, B = P.multiple_zeta(A, add), P.multiple_zeta(B, add)
        for i, b in enumerate(B): A[i] = mul(A[i], b)
        return P.multiple_mobius(A, sub)
    
    def lcm_conv(P, A: list[int], B: list[int], add = operator.add, sub = operator.sub, mul = operator.mul):
        A, B = P.divisor_zeta(A, add), P.divisor_zeta(B, add)
        for i, b in enumerate(B): A[i] = mul(A[i], b)
        return P.divisor_mobius(A, sub)

Back to top page