AoC2021 Day 15

aoc2021_day15.py

import heapq
import numpy as np

x = np.array([[int(c) for c in l.strip()] for l in open('day15.txt')])

def score_map(mp):
    h, w = mp.shape
    def neighbours(y, x):
        for (dy, dx) in [(1, 0), (-1, 0), (0, 1), (0, -1)]:
            if 0 <= y + dy < h and 0 <= x + dx < w:
                yield y + dy, x + dx

    visited = set()
    boundary = [(0, (0, 0))]
    visited.add((0, 0))
    goal = (h-1, w-1)
    score = 0
    while True:
        score, newpos = heapq.heappop(boundary)
        if newpos == goal:
            break
        for n in neighbours(*newpos):
            if n not in visited:
                heapq.heappush(boundary, (score + mp[n], n))
                visited.add(n)
    return score
    
print(f"{score_map(x)}")

h, w = x.shape
big_map = np.zeros([h*5, w*5], dtype=np.int)

for i in range(5):
    big_map[i*h: i*h + h, 0 : w] = x + i
for i in range(5):
    big_map[0: 5*h, w*i : w*i+w] = big_map[0: 5*h, 0:w]+i
big_map = big_map % 9
big_map[big_map==0] = 9

print(f"{score_map(big_map)}")