Make solution for 401 faster

This commit is contained in:
2026-10-03 13:36:16 -04:00
parent 2b0d306081
commit ccc010ccb8
+1 -20
View File
@@ -30,32 +30,13 @@ def sum2(a: int, b: int) -> int:
return r 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: def sigma2caps(n: int) -> int:
res = 0 res = 0
mod = 10**9 mod = 10**9
a = 1 a = 1
while a <= n: while a <= n:
d = n // a d = n // a
b = binsearch(a, d, n) b = n // d
s2 = sum2(a, b) % mod s2 = sum2(a, b) % mod
res += (s2 * d) % mod res += (s2 * d) % mod
res %= mod res %= mod