path: root/day25.py

#!/usr/bin/env python

import numpy as np
with open('day25.txt') as data:
    to_num = {
        '.': 0,
        '>': 1,
        'v': 3
    }
    cucumbers = np.stack([list(map(lambda x: to_num[x], list(line.strip()))) for line in data])

def shift(line, char):
    changed = True
    shifted = np.roll(line, 1)
    diff = line - shifted
    if np.count_nonzero(diff == -char) == 0:
        changed = False
    line[diff == -char] = char
    idx = np.nonzero(diff == -char)[0]
    idx -= 1
    line[idx] = 0
    return changed

count = 0
while True:
    count += 1
    changed = False
    for row in cucumbers[:]:
        changed |= shift(row, 1)
    for col in cucumbers.T:
        changed |= shift(col, 3)
    if not changed:
        print(count)
        break