Problem 14: Two missing numbers

from math import isqrt def find_two_missing(numbers: list[int], n: int) -> set[int]: # This solution only requires O(1) additional memory. # sum numbers and squares of numbers x1 = sum(numbers) x2 = sum(x**2 for x in numbers) # expected sums for full list of numbers y1 = n * (n - 1) // 2 # https://www.wolframalpha.com/input?i=sum+i+with+i+from+0+to+n+-+1 y2 = n * (n - 1) * (2 * n - 1) // 6 # https://www.wolframalpha.com/input?i=sum+i*i+with+i+from+0+to+n+-+1 # solve the following two equations for solutions s1 and s2 # x1 + s1 + s2 = y1 # x2 + s1**2 + s2**2 = y2 d = y1 - x1 s = isqrt(2 * (y2 - x2) - d**2) s1 = (d - s) // 2 s2 = (d + s) // 2 return {s1, s2}