Ahora me siento bien con el codigo python

2020/day_3/toboggan_trajectory.py

@@ -1,22 +1,30 @@
+from math import prod
+
 with open('input') as f:
     data = [line.strip() for line in f]
+    w, h= (len(data[0]), len(data))
 
-width, height = (len(data[0]), len(data))
-slopes = ((1, 1), (3, 1), (5, 1), (7, 1), (1, 2))
-
-result = 1
-for slope in slopes:
+def count_trees(slope):
     trees = 0
-    right, down = slope
+    x, y = slope
 
-    while(down < height):
-        if data[down][right%width] == '#':
+    while y < h:
+        if data[y][x%w] == '#':
             trees += 1
-        right += slope[0]
-        down += slope[1]
 
-    result *= trees
-    print(slope, trees)
+        x += slope[0]
+        y += slope[1]
 
-print(result)
+    return trees
+
+def solve_a():
+    return count_trees((3, 1))
+
+def solve_b():
+    slopes = ((1, 1), (3, 1), (5, 1), (7, 1), (1, 2))
+    return prod([count_trees(slope) for slope in slopes])
+
+
+print(solve_a())
+print(solve_b())