day14.py

day14_part1.py

#!/usr/bin/env python3

from collections import defaultdict
from heapq import heappop, heappush


with open("input.txt") as f:
    empty, rocks, walls = defaultdict(list), defaultdict(list), defaultdict(list)
    
    n = 0
    for i, line in enumerate(f):
        n += 1
        for j, char in enumerate(line):
            match char:
                case 'O':
                    heappush(empty[j], i)
                    heappush(rocks[j], i)
                case '.':
                    heappush(empty[j], i)
                case '#':
                    heappush(walls[j], i)
    
    ans = 0
    for i, rock in rocks.items():
        while rock:
            r = heappop(rock)
            while walls[i] and walls[i][0] < r:
                while empty[i] and empty[i][0] < walls[i][0]:
                    heappop(empty[i])
                heappop(walls[i])
            j = heappop(empty[i])
            ans += n - j
    print(ans)