import numpy as np from utils import open_day nums, boards = open_day(4).read().rstrip().split('\n\n', maxsplit=1) nums = np.fromstring(nums, sep=',', dtype=int) numidxs = np.empty(nums.size, dtype=int) numidxs[nums] = np.arange(nums.size) board, _ = boards.split('\n\n', maxsplit=1) board = [list(row.split()) for row in board.split('\n')] boards = np.fromstring(boards.replace('\n', ' '), sep=' ', dtype=int).reshape(-1, len(board), len(board[0])) rboards = numidxs[boards] colmins = np.amin(np.amax(rboards, axis=1), axis=1) rowmins = np.amin(np.amax(rboards, axis=2), axis=1) mins = np.amin(np.stack((colmins, rowmins)), axis=0) sums = np.sum(boards, axis=(1, 2), where=(rboards > np.expand_dims(mins, (-1, -2)))) * nums[mins] print(sums[np.argmin(mins)]) print(sums[np.argmax(mins)])