import math from decimal import * def solution(s): getcontext().prec = 10000 sqrt_2 = Decimal(2).sqrt() return str(beatty_sum(sqrt_2, int(s))) def beatty_sum(r, n): if n == 0: return 0 if n == 1: return 1 floor_r_n = int(math.floor(r*n)) return natural_number_sum(floor_r_n) - 2 * natural_number_sum(floor_r_n - n) - beatty_sum(r, floor_r_n - n) def natural_number_sum(n): return n*(n+1)/2 def ten(n): return 10**n tests = [ ['77', '4208'], ['1', '1'], ['2', '3'], ['3', '7'], ['4', '12'], ['5', '19'], ['6', '27'], ['7', '36'], ['8', '47'], ['9', '59'], ['10', '73'], [ten(100), '0'] ] for test in tests: result = solution(test[0]) print(test[0], result == test[1], result, test[1])