Updated problem_06 in Project Euler

project_euler/problem_06/sol1.py

@@ -15,7 +15,7 @@
 """
 
 
-def solution(n):
+def solution(n: int) -> int:
     """Returns the difference between the sum of the squares of the first n
     natural numbers and the square of the sum.
 
@@ -28,14 +28,16 @@ def solution(n):
     >>> solution(50)
     1582700
     """
-    suma = 0
-    sumb = 0
+    sum_of_squares = 0
+    sum_of_ints = 0
     for i in range(1, n + 1):
-        suma += i ** 2
-        sumb += i
-    sum = sumb ** 2 - suma
-    return sum
+        sum_of_squares += i ** 2
+        sum_of_ints += i
+    return sum_of_ints ** 2 - sum_of_squares
 
 
 if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()
     print(solution(int(input().strip())))

project_euler/problem_06/sol2.py

@@ -15,7 +15,7 @@
 """
 
 
-def solution(n):
+def solution(n: int) -> int:
     """Returns the difference between the sum of the squares of the first n
     natural numbers and the square of the sum.
 
@@ -28,11 +28,13 @@ def solution(n):
     >>> solution(50)
     1582700
     """
-    suma = n * (n + 1) / 2
-    suma **= 2
-    sumb = n * (n + 1) * (2 * n + 1) / 6
-    return int(suma - sumb)
+    sum_cubes = (n * (n + 1) // 2) ** 2
+    sum_squares = n * (n + 1) * (2 * n + 1) // 6
+    return sum_cubes - sum_squares
 
 
 if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()
     print(solution(int(input().strip())))

project_euler/problem_06/sol3.py

@@ -16,7 +16,7 @@
 import math
 
 
-def solution(n):
+def solution(n: int) -> int:
     """Returns the difference between the sum of the squares of the first n
     natural numbers and the square of the sum.
 
@@ -35,4 +35,7 @@ def solution(n):
 
 
 if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()
     print(solution(int(input().strip())))

project_euler/problem_06/sol4.py

@@ -15,7 +15,7 @@
 """
 
 
-def solution(n):
+def solution(n: int) -> int:
     """Returns the difference between the sum of the squares of the first n
     natural numbers and the square of the sum.
 
@@ -36,4 +36,7 @@ def solution(n):
 
 
 if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()
     print(solution(int(input("Enter a number: ").strip())))

project_euler/problem_06/test_solutions.py

@@ -0,0 +1,18 @@
+from .sol1 import solution as sol1
+from .sol2 import solution as sol2
+from .sol3 import solution as sol3
+from .sol4 import solution as sol4
+
+
+def test_solutions() -> None:
+    """
+    >>> test_solutions()
+    """
+    assert sol1(10) == sol2(10) == sol3(10) == sol4(10) == 2640
+    assert sol1(15) == sol2(15) == sol3(15) == sol4(15) == 13160
+    assert sol1(20) == sol2(20) == sol3(20) == sol4(20) == 41230
+    assert sol1(50) == sol2(50) == sol3(50) == sol4(50) == 1582700
+
+
+if __name__ == "__main__":
+    test_solutions()