From 2de2267319c2c44e416b840daa370ef463d655fe Mon Sep 17 00:00:00 2001
From: Du Yuanchao <shellhub.me@gmail.com>
Date: Thu, 17 Sep 2020 17:37:53 +0800
Subject: [PATCH] Updated problem_06 in Project Euler (#2439)

* * rename variable
* fix type hint
* fix doctest

* added test function

* fixed import error

* fixup! Format Python code with psf/black push

Co-authored-by: github-actions <${GITHUB_ACTOR}@users.noreply.github.com>
---
 project_euler/problem_06/sol1.py           | 16 +++++++++-------
 project_euler/problem_06/sol2.py           | 12 +++++++-----
 project_euler/problem_06/sol3.py           |  5 ++++-
 project_euler/problem_06/sol4.py           |  5 ++++-
 project_euler/problem_06/test_solutions.py | 18 ++++++++++++++++++
 5 files changed, 42 insertions(+), 14 deletions(-)
 create mode 100644 project_euler/problem_06/test_solutions.py

diff --git a/project_euler/problem_06/sol1.py b/project_euler/problem_06/sol1.py
index 513b35467..d6506ea28 100644
--- a/project_euler/problem_06/sol1.py
+++ b/project_euler/problem_06/sol1.py
@@ -15,7 +15,7 @@ numbers and the square of the sum.
 """
 
 
-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())))
diff --git a/project_euler/problem_06/sol2.py b/project_euler/problem_06/sol2.py
index 18cdb5175..5032775f9 100644
--- a/project_euler/problem_06/sol2.py
+++ b/project_euler/problem_06/sol2.py
@@ -15,7 +15,7 @@ numbers and the square of the sum.
 """
 
 
-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())))
diff --git a/project_euler/problem_06/sol3.py b/project_euler/problem_06/sol3.py
index ee739c9a1..385ad0515 100644
--- a/project_euler/problem_06/sol3.py
+++ b/project_euler/problem_06/sol3.py
@@ -16,7 +16,7 @@ numbers and the square of the sum.
 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())))
diff --git a/project_euler/problem_06/sol4.py b/project_euler/problem_06/sol4.py
index 07eed57ba..912a3e8f6 100644
--- a/project_euler/problem_06/sol4.py
+++ b/project_euler/problem_06/sol4.py
@@ -15,7 +15,7 @@ numbers and the square of the sum.
 """
 
 
-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())))
diff --git a/project_euler/problem_06/test_solutions.py b/project_euler/problem_06/test_solutions.py
new file mode 100644
index 000000000..6d5337789
--- /dev/null
+++ b/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()