153. 研究高斯整数(Investigating Gaussian Integers)

#!/usr/bin/env python3
"""Project Euler Problem 153: Investigating Gaussian Integers."""

from __future__ import annotations

import argparse
import math
import sys
import time
import unittest


DEFAULT_LIMIT = 100_000_000
SAMPLE_LIMIT = 5
SAMPLE_TOTAL = 35
CHECK_LIMIT = 100_000
CHECK_TOTAL = 17_924_657_155

def weighted_floor_sum(limit: int) -> int:
    """Return sum(d * floor(limit / d) for 1 <= d <= limit)."""
    total = 0
    start = 1

    while start <= limit:
        quotient = limit // start
        end = limit // quotient
        total += quotient * (start + end) * (end - start + 1) // 2
        start = end + 1

    return total

def solve(limit: int = DEFAULT_LIMIT) -> int:
    """Return the cumulative sum of Gaussian integer divisors up to limit."""
    if limit < 1:
        return 0

    total = weighted_floor_sum(limit)
    floor_sum_cache = {limit: total}
    squares = [n * n for n in range(math.isqrt(limit) + 1)]
    max_small_coordinate = math.isqrt(limit // 2)
    gcd = math.gcd
    cache_get = floor_sum_cache.get

    for x in range(1, max_small_coordinate + 1):
        x_square = squares[x]
        max_y = math.isqrt(limit - x_square)

        if x == 1:
            for y in range(1, max_y + 1):
                radius_square = 1 + squares[y]
                quotient = limit // radius_square
                floor_sum = cache_get(quotient)
                if floor_sum is None:
                    floor_sum = weighted_floor_sum(quotient)
                    floor_sum_cache[quotient] = floor_sum

                coefficient = 2 if y == 1 else 2 * (1 + y)
                total += coefficient * floor_sum
            continue

        start_y = x + 1
        step = 1
        if x % 2 == 0:
            step = 2
            if start_y % 2 == 0:
                start_y += 1

        for y in range(start_y, max_y + 1, step):
            if gcd(x, y) != 1:
                continue

            radius_square = x_square + squares[y]
            quotient = limit // radius_square
            floor_sum = cache_get(quotient)
            if floor_sum is None:
                floor_sum = weighted_floor_sum(quotient)
                floor_sum_cache[quotient] = floor_sum

            total += 2 * (x + y) * floor_sum

    return total

def is_gaussian_divisor(n: int, real: int, imaginary: int) -> bool:
    """Return whether real + imaginary*i divides the rational integer n."""
    norm = real * real + imaginary * imaginary
    return (n * real) % norm == 0 and (n * imaginary) % norm == 0


def brute_force_solve(limit: int) -> int:
    """Return the cumulative divisor sum by direct Gaussian division checks."""
    total = 0

    for n in range(1, limit + 1):
        for real in range(1, n + 1):
            for imaginary in range(-n, n + 1):
                if is_gaussian_divisor(n, real, imaginary):
                    total += real

    return total

class GaussianDivisorTests(unittest.TestCase):
    """Regression tests for the published examples."""

    def test_weighted_floor_sum_for_rational_divisors(self) -> None:
        """Check the ordinary integer divisor contribution for a small limit."""
        self.assertEqual(weighted_floor_sum(5), 21)

    def test_problem_sample(self) -> None:
        """Check the sum of s(n) for the first five positive integers."""
        self.assertEqual(solve(SAMPLE_LIMIT), SAMPLE_TOTAL)

    def test_published_check_value(self) -> None:
        """Check the published value for the first one hundred thousand integers."""
        self.assertEqual(solve(CHECK_LIMIT), CHECK_TOTAL)

    def test_small_limits_against_direct_gaussian_division(self) -> None:
        """Check several small limits against an independent brute-force method."""
        for limit in range(1, 16):
            with self.subTest(limit=limit):
                self.assertEqual(solve(limit), brute_force_solve(limit))

def run_tests() -> None:
    """Run the built-in regression tests."""
    suite = unittest.defaultTestLoader.loadTestsFromTestCase(GaussianDivisorTests)
    result = unittest.TextTestRunner(verbosity=2).run(suite)
    if not result.wasSuccessful():
        raise SystemExit(1)


def main() -> None:
    """Parse command line arguments and print the computed answer."""
    parser = argparse.ArgumentParser(description="Solve Project Euler problem 153.")
    parser.add_argument("--limit", type=int, default=DEFAULT_LIMIT)
    parser.add_argument("--test", action="store_true")
    args = parser.parse_args()

    if args.test:
        run_tests()
        return

    start = time.perf_counter()
    answer = solve(args.limit)
    elapsed = time.perf_counter() - start
    print(f"Result: {answer}")
    print(f"Time: {elapsed:.3f}s")


if __name__ == "__main__":
    main()