Solution for day 15 AoC 2021

day15.py

def main():
    import networkx as nx
    import numpy as np
    data = np.loadtxt("/home/yassinealouini/Documents/code/advent_of_code/aoc/year_2021/data/15.txt", "U")
    size = int(len(data[0])) 
    data = data.view('U1').astype(int).reshape(data.shape[0], size)
    # Need digraph since moving from previous to next position is not the same as the reverse.
    G = nx.DiGraph()

    def new_grid(data):
        # Stack on both axes to build the new grid.
        return np.vstack([
            np.hstack([
                # Back to 1 if 10, 2 if 11, and so on.
                np.where((data + i + j) > 9, (data + i + j) - 9, (data + i + j))
                for i in range(5)
            ])
            for j in range(5)
        ])
    def solve(data, part_2=False):
        data = new_grid(data) if part_2 else data
        npad = ((1, 1), (1, 1))
        data = np.pad(data, pad_width=npad, mode='constant', constant_values=1000)
        grid_size = data.shape[0]
        for x in range(1, grid_size-1):
            for y in range(1, grid_size-1):
                for n in [(x, y+1), (x+1, y), (x-1, y), (x, y-1)]:
                    G.add_edge((x, y), n, weight=data[n])
        return nx.dijkstra_path_length(G, (1, 1), (grid_size-2, grid_size-2), weight='weight')

    print("Solution to part I", solve(data))
    print("Solution to part II", solve(data, part_2=True))

main()