#!/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