day20_part1.py (4671B)
1 #!/usr/bin/env python 2 """Advent of Code 2021, day 20 (part 1): Trench Map 3 Iteratively apply an 'image enhancement algorithm' to a map of the trench 4 floor""" 5 6 # I peeked at the puzzle input to see if I wanted to use sparse or dense 7 # matrices for this, and it looks like dense (i.e., 2d numpy arrays) is the way 8 # to go. 9 # 10 # For the record, I made the (apparently very common) mistake of assuming that 11 # the "world" remained 0 across all enhancements, and did submit an incorrect 12 # answer before I figured it out. 13 14 from typing import Tuple 15 from functools import reduce 16 import numpy as np 17 from utils import get_puzzle_input 18 19 EXAMPLE_INPUT = \ 20 "..#.#..#####.#.#.#.###.##.....###.##.#..###.####..#####..#....#..#..##..##" + \ 21 "#..######.###...####..#..#####..##..#.#####...##.#.#..#.##..#.#......#.###" + \ 22 ".######.###.####...#.##.##..#..#..#####.....#.#....###..#.##......#.....#." + \ 23 ".#..#..##..#...##.######.####.####.#.#...#.......#..#.#.#...####.##.#....." + \ 24 ".#..#...##.#.##..#...##.#.##..###.#......#.#.......#.#.#.####.###.##...#.." + \ 25 "...####.#..#..#.##.#....##..#.####....##...##..#...#......#.#.......#....." + \ 26 """..##..####..#...#.#.#...##..#.#..###..#####........#..####......#..# 27 28 #..#. 29 #.... 30 ##..# 31 ..#.. 32 ..### 33 """ 34 35 def parse_input(input_string: str) -> Tuple[np.ndarray, np.ndarray]: 36 """Given the puzzle input, return a tuple containing the image data (a 2d 37 numpy array of integers) and a lookup table (a 1d numpy array of integers)""" 38 # I'm ditching the '#' / '.' representation right away; we want zeroes and 39 # ones here 40 lines = input_string.rstrip('\n').split('\n') 41 lookup_table = np.array([{'#': 1, '.': 0}[c] for c in lines[0]]) 42 image_data = np.array( 43 [[{'#': 1, '.': 0}[c] for c in l] for l in lines[2:]] 44 ) 45 return (image_data, lookup_table) 46 47 def dilate_image(image_data: np.ndarray, 48 dilate_with: int) -> np.ndarray: 49 """Given an n x m image matrix, return a (n + 2) x (m + 2) matrix of the 50 same image padded around with `dilate_with`""" 51 new_image = np.zeros((image_data.shape[0] + 2, image_data.shape[1] + 2), 52 dtype=int) + dilate_with 53 new_image[1:-1, 1:-1] = image_data[:, :] 54 return new_image 55 56 def find_pixel_indices(image_data: np.ndarray, 57 world_state: int) -> np.ndarray: 58 """Use each pixel's neighborhood to convert it to an index (in the range 59 0..512); given an n x m image returns an (n + 2) x (m + 2) set of indices""" 60 # Using dilate this way is "inefficient", but it's fast enough that the I 61 # think it's a good ease of programming / speed of execution tradeoff. 62 image_dilated = dilate_image(dilate_image(image_data, world_state), 63 world_state) 64 image_indices = dilate_image(np.zeros_like(image_data), 0) 65 for exponent in range(9): 66 col_offset = 2 - exponent % 3 67 row_offset = 2 - exponent // 3 68 image_indices += 2**exponent * image_dilated[ 69 row_offset:row_offset+image_indices.shape[0], 70 col_offset:col_offset+image_indices.shape[1] 71 ] 72 return image_indices 73 74 def enhance_image_once(image_data: np.ndarray, 75 world_state: int, 76 lookup_table: np.ndarray) -> Tuple[np.ndarray, int]: 77 """For each pixel in the image, consider its 8-connected neighborhood; 78 these values represent bits in an integer, which is then used to index into 79 a lookup table to find the new value for the pixel and a new `world_state`""" 80 # World state is either 0 or 1. If it's 0, every unexamined pixel in the 81 # world is replaced with `lookup_table[0]`, and if it's 1 they're replaced 82 # with `lookup_table[511]` 83 return (lookup_table[find_pixel_indices(image_data, world_state)], 84 int(lookup_table[{0: 0b0, 1: 0b111111111}[world_state]])) 85 86 def enhance_image(image_data: np.ndarray, lookup_table: np.ndarray, 87 num_enhancements: int) -> Tuple[np.ndarray, int]: 88 """Apply the enhancement `num_enhancements` times, starting with an initial 89 `world_state` of 0""" 90 return reduce(lambda x, y: enhance_image_once(x[0], x[1], lookup_table), 91 range(num_enhancements), 92 (image_data, 0)) 93 94 def solve_puzzle(input_string: str) -> int: 95 """Return the numeric solution to the puzzle""" 96 return enhance_image(*((parse_input(input_string)) + (2,)))[0].sum() 97 98 def main() -> None: 99 """Run when the file is called as a script""" 100 assert solve_puzzle(EXAMPLE_INPUT) == 35 101 print("Total number of lit pixels:", 102 solve_puzzle(get_puzzle_input(20))) 103 104 if __name__ == "__main__": 105 main()