diff --git a/src/Python/Problem061.py b/src/Python/Problem061.py
index ef648ed..a68ad54 100644
--- a/src/Python/Problem061.py
+++ b/src/Python/Problem061.py
@@ -9,6 +9,8 @@ Solution for problem 61 of Project Euler
 https://projecteuler.net/problem=61
 """
 
+from itertools import permutations
+
 from utils import timeit
 
 
@@ -19,7 +21,7 @@ def compute():
     are all figurate (polygonal) numbers and are generated by the following
     formulae:
         Triangle 	  	P3,n=n(n+1)/2 	  	1, 3, 6, 10, 15, ...
-        Square 	  	P4,n=n2 	  	1, 4, 9, 16, 25, ...
+        Square 	  	    P4,n=n2 	  	    1, 4, 9, 16, 25, ...
         Pentagonal 	  	P5,n=n(3n-1)/2 	  	1, 5, 12, 22, 35, ...
         Hexagonal 	  	P6,n=n(2n-1) 	  	1, 6, 15, 28, 45, ...
         Heptagonal 	  	P7,n=n(5n-3)/2 	  	1, 7, 18, 34, 55, ...
@@ -41,46 +43,33 @@ def compute():
     heptagonal, and octagonal, is represented by a different number in the set.
     """
 
-    from itertools import permutations
+    def F(k, n):
+        return (n * (n - 1) * (k - 2) // 2) + n
 
-    P = {
-        "3": lambda n: n * (n + 1) // 2,
-        "4": lambda n: n * n,
-        "5": lambda n: n * (3 * n - 1) // 2,
-        "6": lambda n: n * (2 * n - 1),
-        "7": lambda n: n * (5 * n - 3) // 2,
-        "8": lambda n: n * (3 * n - 2),
-    }
+    P = {k: set(str(F(k, n)) for n in range(150)) for k in range(3, 9)}
 
-    # Set of values in each figurate set
-    dic = {
-        k: [str(v(i)) for i in range(1, 200) if 10**3 <= v(i) < 10**4]
-        for k, v in P.items()
-    }
+    four_digits_numbers = {str(n) for n in range(10**3, 10**4)}
 
-    perms = permutations(P.keys(), 6)
-    for perm in perms:  # 720 different permutations of the six polygonals
+    dic = {k: list(v & four_digits_numbers) for k, v in P.items()}
+
+    perms = permutations(dic.keys(), 6)
+    for perm in perms:
         for a in dic.get(perm[0]):
             lst = []
             a1, a2 = a[:2], a[2:]
-            lst.append(a)
             for b in filter(lambda n: n[:2] == a2, dic.get(perm[1])):
-                lst.append(b)
                 b2 = b[2:]
                 for c in filter(lambda n: n[:2] == b2, dic.get(perm[2])):
-                    lst.append(c)
                     c2 = c[2:]
                     for d in filter(lambda n: n[:2] == c2, dic.get(perm[3])):
-                        lst.append(d)
                         d2 = d[2:]
                         for e in filter(lambda n: n[:2] == d2, dic.get(perm[4])):
-                            lst.append(e)
                             e2 = e[2:]
                             for f in filter(
                                 lambda n: n[:2] == e2 and n[2:] == a1, dic.get(perm[5])
                             ):
-                                lst.append(f)
-                                return sum(map(int, lst))
+                                lst.append([a, b, c, d, e, f])
+                                return sum(map(int, lst[0]))
 
 
 if __name__ == "__main__":