cp_library/math/table/sieve_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 functools import cached_property
from cp_library.math.table.sieve_proto import SieveProtocol
from cp_library.math.table.primes_cls import Primes

class Sieve(list[int], SieveProtocol):
    def __init__(spf, N):
        super().__init__(i for i in range(N+1))
        spf[0] = 1
        for x in range(2, N+1):
            x2 = x*x
            if x2 > N: break
            if spf[x] == x:
                for j in range(x2, N+1, x):
                    if spf[j] == j:
                        spf[j] = x
    @cached_property
    def primes(spf) -> Primes:
        gen = (x for x,f in enumerate(spf) if f == x)
        primes = Primes.__new__(Primes)
        super(Primes, primes).__init__(gen)
        return primes

    def factor_cnts(spf, N):
        assert N < len(spf)
        pairs = []
        while N > 1:
            match pairs:
                case [*_, (f,cnt)] if f == spf[N]:
                    pairs[-1] = (f,cnt+1)
                case _:
                    pairs.append((spf[N], 1))
            N //= spf[N]
        return pairs

    def factors(spf, N):
        assert N < len(spf)
        factors = []
        while N > 1:
            factors.append(spf[N])
            N //= spf[N]
        return factors
    
    def unique_factors(spf, N):
        assert N < len(spf)
        factors = []
        while N > 1:
            if factors and factors[-1] != spf[N]: 
                factors.append(spf[N])
            N //= spf[N]
        return factors

'''
╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸
             https://kobejean.github.io/cp-library               
'''
from functools import cached_property
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 Sieve(list[int], SieveProtocol):
    def __init__(spf, N):
        super().__init__(i for i in range(N+1))
        spf[0] = 1
        for x in range(2, N+1):
            x2 = x*x
            if x2 > N: break
            if spf[x] == x:
                for j in range(x2, N+1, x):
                    if spf[j] == j:
                        spf[j] = x
    @cached_property
    def primes(spf) -> Primes:
        gen = (x for x,f in enumerate(spf) if f == x)
        primes = Primes.__new__(Primes)
        super(Primes, primes).__init__(gen)
        return primes

    def factor_cnts(spf, N):
        assert N < len(spf)
        pairs = []
        while N > 1:
            match pairs:
                case [*_, (f,cnt)] if f == spf[N]:
                    pairs[-1] = (f,cnt+1)
                case _:
                    pairs.append((spf[N], 1))
            N //= spf[N]
        return pairs

    def factors(spf, N):
        assert N < len(spf)
        factors = []
        while N > 1:
            factors.append(spf[N])
            N //= spf[N]
        return factors
    
    def unique_factors(spf, N):
        assert N < len(spf)
        factors = []
        while N > 1:
            if factors and factors[-1] != spf[N]: 
                factors.append(spf[N])
            N //= spf[N]
        return factors

Back to top page