Problem 14: Two missing numbers

def find_two_missing(numbers: list[int], n: int) -> set[int]: # Sum of full range from 0..n-1 full_sum = (n * (n - 1)) // 2 given_sum = sum(numbers) missing_sum = full_sum - given_sum # a + b # Sum of squares of full range full_sq = (n - 1) * n * (2 * n - 1) // 6 given_sq = sum(x * x for x in numbers) missing_sq = full_sq - given_sq # a² + b² # a*b = [(a+b)² - (a²+b²)] / 2 prod = (missing_sum * missing_sum - missing_sq) // 2 # Solve quadratic: t² - (a+b)t + ab == 0 # discriminant = (a + b)² - 4ab D = missing_sum * missing_sum - 4 * prod root_D = int(D ** 0.5) while (root_D + 1) * (root_D + 1) <= D: root_D += 1 while root_D * root_D > D: root_D -= 1 a = (missing_sum - root_D) // 2 b = missing_sum - a return {a, b} # quick sanity checks if __name__ == "__main__": import random assert find_two_missing([], 2) == {0, 1} assert find_two_missing([0], 3) == {1, 2} assert find_two_missing([1], 3) == {0, 2} assert find_two_missing([2], 3) == {0, 1} assert find_two_missing([1, 2], 4) == {0, 3} assert find_two_missing([4, 2, 0], 5) == {1, 3} assert find_two_missing([0, 1, 2, 5], 6) == {3, 4} assert find_two_missing([5, 2, 1, 0], 6) == {3, 4} def test(): for n in range(2, 101): numbers = list(range(n)) random.shuffle(numbers) real_missing_numbers = {numbers.pop() for _ in range(2)} random.shuffle(numbers) missing_numbers = find_two_missing(numbers, n) assert real_missing_numbers == missing_numbers, f"Test failed for {n=}. Expected {real_missing_numbers}, got {missing_numbers} for {numbers=}" test() print("All tests passed!")