Update sol4.py

Author: MaximSmolskiy

project_euler/problem_009/sol4.py

@@ -5,23 +5,32 @@
 
 A Pythagorean triplet is a set of three natural numbers, a < b < c, for which,
 
-    a^2 + b^2 = c^2
+    a^2 + b^2 = c^2.
 
 For example, 3^2 + 4^2 = 9 + 16 = 25 = 5^2.
 
 There exists exactly one Pythagorean triplet for which a + b + c = 1000.
-Find the product a*b*c.
+Find the product abc.
 
 References:
     - https://en.wikipedia.org/wiki/Pythagorean_triple
 """
 
 
 def get_squares(n: int) -> list[int]:
-    res = [0] * n
-    for i in range(n):
-        res[i] = i * i
-    return res
+    """
+    >>> get_squares(0)
+    []
+    >>> get_squares(1)
+    [0]
+    >>> get_squares(2)
+    [0, 1]
+    >>> get_squares(3)
+    [0, 1, 4]
+    >>> get_squares(4)
+    [0, 1, 4, 9]
+    """
+    return [number * number for number in range(n)]
 
 
 def solution(n: int = 1000) -> int: