advent_of_code_2021

My attempts to work through the 2021 Advent of Code problems.
git clone https://git.eamoncaddigan.net/advent_of_code_2021.git
Log | Files | Refs | README | LICENSE

day25_part1.py (3835B)

      1 #!/usr/bin/env python
      2 """Advent of Code 2021, day 25 (part 1): Sea Cucumber
      3 Simulate the movement of sea cucumbers"""
      4 
      5 # Going to use NumPy again. I'm going to represent the seafloor grid using
      6 # integers, with 0 corresponding to an empty space, 1 corresponding to an
      7 # east-moving sea cucumber, and 2 corresponding to a south-moving sea cucumber.
      8 # E.g., the puzzle input:
      9 #
     10 # v...>>.vv>
     11 # .vv>>.vv..
     12 # >>.>v>...v
     13 # >>v>>.>.v.
     14 # v>v.vv.v..
     15 # >.>>..v...
     16 # .vv..>.>v.
     17 # v.v..>>v.v
     18 # ....v..v.>
     19 #
     20 # Becomes:
     21 #
     22 # 2000110221
     23 # 0221102200
     24 # 1101210002
     25 # 1121101020
     26 # 2120220200
     27 # 1011002000
     28 # 0220010120
     29 # 2020011202
     30 # 0000200201
     31 
     32 #from typing import Tuple, List, NewType
     33 
     34 import numpy as np
     35 
     36 from utils import convert_input_to_array, get_puzzle_input
     37 
     38 EXAMPLE_INPUT = \
     39 """v...>>.vv>
     40 .vv>>.vv..
     41 >>.>v>...v
     42 >>v>>.>.v.
     43 v>v.vv.v..
     44 >.>>..v...
     45 .vv..>.>v.
     46 v.v..>>v.v
     47 ....v..v.>
     48 """
     49 
     50 def parse_input(input_string: str) -> np.ndarray:
     51     """Given the puzzle input, return a 2d numpy array"""
     52     return convert_input_to_array(
     53         input_string
     54         .replace('.', '0')
     55         .replace('>', '1')
     56         .replace('v', '2')
     57     )
     58 
     59 def shift_up(cucumber_map: np.ndarray) -> np.ndarray:
     60     """Return the entire map shifted down"""
     61     return np.vstack((cucumber_map[1:, :], cucumber_map[0, :]))
     62 
     63 def shift_down(cucumber_map: np.ndarray) -> np.ndarray:
     64     """Return the entire map shifted down"""
     65     return np.vstack((cucumber_map[-1, :], cucumber_map[0:-1, :]))
     66 
     67 def is_empty(cucumber_map: np.ndarray, direction: str) -> np.ndarray:
     68     """Given a cucumber map, check if the adjacent space in the given direction
     69     is empty"""
     70     if direction == 'east':
     71         return np.transpose(is_empty(np.transpose(cucumber_map), 'south'))
     72     # So actually we just need to check that the space to the south is empty
     73     # (remembering that we wrap around to the top row from the bottom
     74     assert direction == 'south'
     75     return shift_up(cucumber_map) == 0
     76 
     77 def shift_selected(cucumber_map: np.ndarray, movers: np.ndarray) -> np.ndarray:
     78     """Shift the selected elements down one row"""
     79     new_map = cucumber_map.copy()
     80     new_map[shift_down(movers).nonzero()] = cucumber_map[movers.nonzero()]
     81     new_map[movers.nonzero()] = 0
     82     return new_map
     83 
     84 def move_herd(cucumber_map: np.ndarray, direction: str) -> np.ndarray:
     85     """Move the herd in the given direction"""
     86     movers = is_empty(cucumber_map, direction)
     87     if direction == 'east':
     88         movers &= cucumber_map == 1
     89         new_map = np.transpose(shift_selected(np.transpose(cucumber_map),
     90                                               np.transpose(movers)))
     91     elif direction == 'south':
     92         movers &= cucumber_map == 2
     93         new_map = shift_selected(cucumber_map, movers)
     94     else:
     95         raise ValueError('Direction must be "east" or "south"')
     96     return new_map
     97 
     98 def step_herd(cucumber_map: np.ndarray) -> np.ndarray:
     99     """Move the whole herd east, and then south"""
    100     return move_herd(move_herd(cucumber_map, 'east'), 'south')
    101 
    102 def run_until_stuck(cucumber_map: np.ndarray) -> int:
    103     """Keep stepping the herd through its moves until it stops moving"""
    104     herd_position = cucumber_map
    105     for step_number in range(1_000_000):
    106         new_position = step_herd(herd_position)
    107         if (new_position == herd_position).all():
    108             return step_number+1
    109         herd_position = new_position
    110     raise RuntimeError('Failed to lock up after very many steps')
    111 
    112 def solve_puzzle(input_string: str) -> int:
    113     """Return the numeric solution to the puzzle"""
    114     return run_until_stuck(parse_input(input_string))
    115 
    116 def main() -> None:
    117     """Run when the file is called as a script"""
    118     assert solve_puzzle(EXAMPLE_INPUT) == 58
    119     print("Steps for cucumbers to become stuck:",
    120           solve_puzzle(get_puzzle_input(25)))
    121 
    122 if __name__ == "__main__":
    123     main()