aoc_day3_2020.py

WIDTH = len(INPUT.split("\n")[0])

# Part 1
trees = 0
for i, line in enumerate(INPUT.split("\n")):
    if line[i * 3 % WIDTH] == '#':
        trees += 1

print(trees)

WIDTH = len(INPUT.split("\n")[0])

# Part 2
product = 1
for right_shift, down_shift in [(1, 1), (3, 1), (5, 1), (7, 1), (1, 2)]:
    trees = 0
    for i, line in enumerate(INPUT.split("\n")):
        if i % down_shift == 0:
            if line[int(i * (right_shift / down_shift)) % WIDTH] == '#':
                trees += 1
    product *= trees

print(product)