Solve problem 401

python/e401.py

@@ -0,0 +1,78 @@
+def divisors(n: int) -> list[int]:
+    xs = set()
+    d = 1
+    while d * d <= n:
+        if n % d == 0:
+            xs.add(d)
+            xs.add(n // d)
+        d += 1
+    return sorted(list(xs))
+
+
+def sigma2(n: int) -> int:
+    return sum(d**2 for d in divisors(n))
+
+
+def sum2(a: int, b: int) -> int:
+    return sum(x**2 for x in range(a, b + 1))
+
+
+def sigma2caps_naive(n: int) -> int:
+    return sum(sigma2(i) for i in range(1, n + 1))
+
+
+def sum2(a: int, b: int) -> int:
+    # Calculate sum(x**2 for x in range(a, b + 1)) directly
+    # https://stackoverflow.com/questions/70902823/
+    r = b * (b + 1) * (2 * b + 1)
+    r -= (a - 1) * a * (2 * (a - 1) + 1)
+    r //= 6
+    return r
+
+
+def binsearch(a: int, d: int, n: int) -> int:
+    # Return the largest integer i in [a, n] for which n // i == d.
+
+    # Naive:
+    # b = a
+    # while n // (b + 1) == d:
+    #     b += 1
+    # return b
+
+    b = n
+    while n // b != d:
+        mid = (a + b) // 2
+        if b // mid != d:
+            b = mid
+        else:
+            a = mid
+    return b
+
+
+def sigma2caps(n: int) -> int:
+    res = 0
+    mod = 10**9
+    a = 1
+    while a <= n:
+        d = n // a
+        b = binsearch(a, d, n)
+        s2 = sum2(a, b) % mod
+        res += (s2 * d) % mod
+        res %= mod
+        a = b + 1
+    return res
+
+
+def euler_401():
+    assert sum2(23, 1337) == sum(x**2 for x in range(23, 1337 + 1))
+    assert sigma2caps_naive(6) == 113
+    assert sigma2caps(6) == 113
+    assert sigma2caps(10**6) == 385965077
+    assert sigma2caps(10**9) == 73475174
+    return sigma2caps(10**15)
+
+
+if __name__ == "__main__":
+    solution = euler_401()
+    print(f"e401.py: {solution}")
+    assert solution == 281632621

python/lib_create_templates.py

@@ -8,7 +8,7 @@ def euler_XXX():
 if __name__ == "__main__":
     solution = euler_XXX()
     print(f"eXXX.py: {solution}")
-    # assert(solution == 0)
+    # assert solution == 0
 
 """