diff --git a/2020/day01-expense.py b/2020/day01-expense.py
index 471d5fd..0c9fd74 100644
--- a/2020/day01-expense.py
+++ b/2020/day01-expense.py
@@ -1,18 +1,23 @@
+from itertools import combinations
+from math import prod
+
+
 def main():
     lines = []
     with open("inputs/day01") as f:
         for line in f:
             lines.append(int(line.strip()))
 
-    stop = len(lines)
-    for i, value in enumerate(lines):
-        for j in range(i+1, stop):
-            other = lines[j]
-            for k in range(j+1, stop):
-                third = lines[k]
-                addition = other + value + third
-                if addition == 2020:
-                    print("result is", other * value * third)
+    res = solve(lines, 2)
+    print("result is", res)
+    res = solve(lines, 3)
+    print("result is", res)
+
+
+def solve(expense_report, fix_number):
+    for combination in combinations(expense_report, fix_number):
+        if sum(combination) == 2020:
+            return prod(combination)
 
 
 if __name__ == '__main__':