use beatty sum

Verifying solution...
Test 1 passed!
Test 2 passed!
Test 3 passed! [Hidden]
Test 4 failed  [Hidden]
Test 5 failed  [Hidden]
Test 6 failed  [Hidden]
Test 7 failed  [Hidden]
Test 8 failed  [Hidden]
Test 9 failed  [Hidden]
Test 10 failed  [Hidden]

others/dodge-the-lasers/solution.py

@@ -1,49 +1,26 @@
 import math
 from decimal import *
 
-getcontext().prec = 102
-sqrt_2 = Decimal(2).sqrt()
-
 def solution(s):
-    total = 0
+    getcontext().prec = 10000
-    for i in range(int(s)+1):
+    sqrt_2 = Decimal(2).sqrt()
-        total += int(math.floor(i*sqrt_2))
-        #print i, total
-    return str(total)
 
-def check_known_pattern(start, end, memo):
+    return str(beatty_sum(sqrt_2, int(s)))
-    if '0'*(end-start+1) in memo: # TODO: checkk value in the memo
-        return '0'*(end-start+1)
-    return False
 
-def generate_dfs(start, end, memo):
+def beatty_sum(r, n):
-    pattern = check_known_pattern(start, end, memo)
+    if n == 0:
-    if pattern:
+        return 0
-        yield pattern, memo[pattern]
+    if n == 1:
-    else:
+        return 1
-        new_pattern = ''
+    floor_r_n = int(math.floor(r*n))
-        new_value = 0 # TODO: calculate new_value
-        pivot = int(math.floor((start+end)/2)) # TODO: power of 2 is optimal
-        for pattern, value in generate_dfs(start, pivot, memo):
-            new_pattern += pattern
-            new_value += value
-            yield pattern, value
-        for pattern, value in generate_dfs(pivot + 1, end, memo):
-            new_pattern += pattern
-            new_value += value
-            yield pattern, value
-        memo[new_pattern] = new_value
     
-memo = {'0': 1, '1': 1} # TODO: upgrade memo
+    return natural_number_sum(floor_r_n) - 2 * natural_number_sum(floor_r_n - n) - beatty_sum(r, floor_r_n - n)
-total = 0
-for pattern, value in generate_dfs(1, 5, memo):
-    total += value
-    print pattern, value
-print total, memo
-
-exit(0)
 
+def natural_number_sum(n):
+    return n*(n+1)/2
 
+def ten(n):
+    return 10**n
 
 tests = [
         ['77', '4208'],
@@ -57,6 +34,7 @@ tests = [
         ['8', '47'],
         ['9', '59'],
         ['10', '73'],
+        [ten(100), '0']
         ]