cp_library/math/table/linear_sieve_cnts_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
• cp_library/math/table/primes_cls.py
• cp_library/math/table/sieve_proto.py

Code

import cp_library.math.table.__header__
from cp_library.math.table.sieve_proto import SieveProtocol
from cp_library.math.table.primes_cls import Primes

class LinearSieveCounts(list[int], SieveProtocol):

    def __init__(spf, N: int):
        super().__init__([0] * (N + 1))
        exp = [0] * (N + 1)
        nxt = [0] * (N + 1)
        primes = Primes.__new__(Primes)
        spf[0], spf[1] = 0, 1
        exp[1] = 1
        for x in range(2,N+1):
            if spf[x] == 0:
                spf[x],exp[x] = x,1
                primes.append(x)
            for p in primes:
                if (y := x*p) > N or p > spf[x]: break
                spf[y] = p
                if x%p:
                    nxt[y], exp[y] = x, 1
                else:
                    nxt[y], exp[y] = nxt[x], exp[x]+1
        spf.primes = primes
        spf.exp = exp
        spf.nxt = nxt
    
    def factor_cnts(spf, N: int):
        assert N < len(spf)
        exp,nxt = spf.exp, spf.nxt
        pairs = []
        while spf[N] != N:
            pairs.append((spf[N],exp[N]))
            N = nxt[N]
        if N:
            pairs.append((spf[N],exp[N]))
        return pairs

    def factors(spf, N):
        assert N < len(spf)
        exp,nxt = spf.exp, spf.nxt
        factors = []
        while N > 1:
            factors.extend(spf[N] for _ in range(exp[N]))
            N = nxt[N]
        return factors

'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
             https://kobejean.github.io/cp-library               
'''
from typing import Protocol
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)

class SieveProtocol(Protocol):
    primes: Primes
    def factor_cnts(self, N): ...
    def factors(self, N): ...
    def unique_factors(self, N): ...
    def __getitem__(self, key) -> int: ...

class LinearSieveCounts(list[int], SieveProtocol):

    def __init__(spf, N: int):
        super().__init__([0] * (N + 1))
        exp = [0] * (N + 1)
        nxt = [0] * (N + 1)
        primes = Primes.__new__(Primes)
        spf[0], spf[1] = 0, 1
        exp[1] = 1
        for x in range(2,N+1):
            if spf[x] == 0:
                spf[x],exp[x] = x,1
                primes.append(x)
            for p in primes:
                if (y := x*p) > N or p > spf[x]: break
                spf[y] = p
                if x%p:
                    nxt[y], exp[y] = x, 1
                else:
                    nxt[y], exp[y] = nxt[x], exp[x]+1
        spf.primes = primes
        spf.exp = exp
        spf.nxt = nxt
    
    def factor_cnts(spf, N: int):
        assert N < len(spf)
        exp,nxt = spf.exp, spf.nxt
        pairs = []
        while spf[N] != N:
            pairs.append((spf[N],exp[N]))
            N = nxt[N]
        if N:
            pairs.append((spf[N],exp[N]))
        return pairs

    def factors(spf, N):
        assert N < len(spf)
        exp,nxt = spf.exp, spf.nxt
        factors = []
        while N > 1:
            factors.extend(spf[N] for _ in range(exp[N]))
            N = nxt[N]
        return factors

Back to top page