540. 统计本原毕达哥拉斯三元组(Counting Primitive Pythagorean Triples)

"""Project Euler 540 - Counting Primitive Pythagorean Triples.

This solver includes an optional Numba acceleration path.
"""

from __future__ import annotations

from functools import lru_cache
from math import isqrt
from time import perf_counter

TARGET_N = 3_141_592_653_589_793
SAMPLE_N = 10**6
SAMPLE_ANSWER = 159_139

try:
    import numpy as np
    from numba import njit

    HAS_NUMBA = True
except Exception:
    HAS_NUMBA = False


if HAS_NUMBA:

    @njit(cache=False)
    def isqrt_numba(n: int) -> int:
        """Return floor(sqrt(n)) for positive integer n in nopython mode."""
        x = int(n**0.5)
        while (x + 1) * (x + 1) <= n:
            x += 1
        while x * x > n:
            x -= 1
        return x


    @njit(cache=False)
    def mobius_sieve_numba(limit: int):
        """Return Mu[0..limit] using an Eratosthenes-style Mobius sieve."""
        mu = np.ones(limit + 1, dtype=np.int8)
        mu[0] = 0
        is_composite = np.zeros(limit + 1, dtype=np.uint8)

        for p in range(2, limit + 1):
            if is_composite[p] != 0:
                continue
            for multiple in range(p, limit + 1, p):
                is_composite[multiple] = 1
                mu[multiple] = -mu[multiple]
            square = p * p
            if square <= limit:
                for multiple in range(square, limit + 1, square):
                    mu[multiple] = 0
        return mu


    @njit(cache=False)
    def build_mertens_prefix_numba(mu):
        """Return prefix sums M(n)=sum_{k<=n} mu(k) for fast grouped summation."""
        mertens = np.zeros(mu.size, dtype=np.int32)
        running = 0
        for i in range(1, mu.size):
            running += int(mu[i])
            mertens[i] = running
        return mertens


    @njit(cache=False)
    def count_ordered_pairs_numba(n: int) -> int:
        """Count pairs (x,y) with 1<=x<y and x^2+y^2<=n."""
        x = 1
        x_sq = 1
        y = isqrt_numba(n)
        y_sq = y * y
        total = 0

        while x < y:
            while x_sq + y_sq > n:
                y_sq -= 2 * y - 1
                y -= 1
            if x >= y:
                break
            total += y - x
            x += 1
            x_sq += 2 * x - 1
        return total


def mobius_sieve(limit: int):
    """Return Mu[0..limit] using an Eratosthenes-style Mobius sieve."""
    if HAS_NUMBA:
        return mobius_sieve_numba(limit)

    mu = [1] * (limit + 1)
    mu[0] = 0
    is_composite = bytearray(limit + 1)

    for p in range(2, limit + 1):
        if is_composite[p]:
            continue
        for multiple in range(p, limit + 1, p):
            is_composite[multiple] = 1
            mu[multiple] = -mu[multiple]
        square = p * p
        if square <= limit:
            for multiple in range(square, limit + 1, square):
                mu[multiple] = 0
    return mu


def build_mertens_prefix(mu):
    """Return prefix sums M(n)=sum_{k<=n} mu(k) for fast grouped summation."""
    if HAS_NUMBA:
        return build_mertens_prefix_numba(mu)

    mertens = [0] * len(mu)
    running = 0
    for i in range(1, len(mu)):
        running += mu[i]
        mertens[i] = running
    return mertens


def count_ordered_pairs(n: int) -> int:
    """Count pairs (x,y) with 1<=x<y and x^2+y^2<=n."""
    if HAS_NUMBA:
        return int(count_ordered_pairs_numba(n))

    x = 1
    x_sq = 1
    y = isqrt(n)
    y_sq = y * y
    total = 0

    while x < y:
        while x_sq + y_sq > n:
            y_sq -= 2 * y - 1
            y -= 1
        if x >= y:
            break
        total += y - x
        x += 1
        x_sq += 2 * x - 1
    return total


def primitive_triples_upto(n: int) -> int:
    """Return P(n): primitive Pythagorean triples with a<b<c<=n."""
    max_d = isqrt(n)
    mu = mobius_sieve(max_d)
    mertens = build_mertens_prefix(mu)

    @lru_cache(maxsize=None)
    def r(value: int) -> int:
        """Memoized wrapper for R(value)."""
        return count_ordered_pairs(value)

    @lru_cache(maxsize=None)
    def q(value: int) -> int:
        """Compute Q(value)=sum mu(d)*R(value//d^2) with grouped d-ranges."""
        total = 0
        d = 1
        limit = isqrt(value)
        while d <= limit:
            quotient = value // (d * d)
            right = isqrt(value // quotient)
            coeff = int(mertens[right]) - int(mertens[d - 1])
            if coeff:
                total += coeff * r(quotient)
            d = right + 1
        return total

    answer = 0
    sign = 1
    factor = 1
    while factor <= n:
        answer += sign * q(n // factor)
        factor <<= 1
        sign = -sign
    return answer


def solve_and_measure(n: int) -> tuple[int, float]:
    """Return (P(n), elapsed_seconds)."""
    start = perf_counter()
    value = primitive_triples_upto(n)
    elapsed = perf_counter() - start
    return value, elapsed


def main() -> None:
    """Run sample verification and print the final answer."""
    assert primitive_triples_upto(20) == 3, "P(20) must be 3"

    sample_value, sample_elapsed = solve_and_measure(SAMPLE_N)
    assert sample_value == SAMPLE_ANSWER, (
        f"P(10^6) mismatch: got {sample_value}, expected {SAMPLE_ANSWER}"
    )
    print(f"P(10^6) = {sample_value}")
    print(f"sample_elapsed_seconds = {sample_elapsed:.6f}")

    final_value, final_elapsed = solve_and_measure(TARGET_N)
    print(f"P({TARGET_N}) = {final_value}")
    print(f"final_elapsed_seconds = {final_elapsed:.6f}")


if __name__ == "__main__":
    main()