day02_part2.py (3475B)
1 #!/usr/bin/env python 2 """Day 2, Part 2: It turns out that controlling the submarine is more 3 complicated than we thought in Part 1. We need to maintain a third state 4 variable (in addition to horizontal and vertical position) called 'aim', and 5 use that when calculating sub movement.""" 6 7 # Hey, the functions from Part 1 will be useful once again. All the input 8 # parsing is the same, and now we just need to use slightly more complicated 9 # logic when figuring out final position. We can start from the point where we 10 # have a data frame with the columns `axis` and `offset`. 11 12 #from io import StringIO 13 import numpy as np 14 import pandas as pd 15 from utils import get_puzzle_input 16 from day02_part1 import (convert_input_to_df, 17 split_command_parts, 18 convert_commands_to_offsets) 19 20 def find_submarine_state(offset_df): 21 """Given a pandas data frame with the columns `axis` and `offset`, we track 22 `aim` as well as `horizontal` and `vertical` offsets, using complicated 23 rules: 24 * 'vertical X' increases or decreases aim (which starts at 0) 25 * 'horizontal X': 26 * increases horizontal position by X 27 * increases vertical position (i.e., depth) by aim * X 28 This returns a data frame with the columns `horizontal`, `vertical` and 29 `aim`, representing the running value of these states.""" 30 return ( 31 pd.DataFrame( 32 # 'aim' represents the up-to-date state of the aim parameter, while 33 # 'horizontal_step' is just the current step's offset. We save this 34 # intermediate calculation because we need it for 'depth'. 35 # This line uses the extremely new-to-me 'walrus operator' to 36 # cretae the variable `vert_rows` so we don't calculate the 37 # equality twice. Not sure how I feel about this style. 38 {'aim': np.where(vert_rows := offset_df['axis'] == 'vertical', 39 offset_df['offset'], 0).cumsum(), 40 'horizontal_step': np.where(~vert_rows, 41 offset_df['offset'], 0)} 42 ) 43 .assign(depth=lambda x: (x['aim'] * x['horizontal_step']).cumsum(), 44 horizontal=lambda x: x['horizontal_step'].cumsum()) 45 .drop('horizontal_step', axis='columns') 46 ) 47 48 def calculate_puzzle_solution(running_state_df): 49 """Given a pandas data frame with the running values of the state of the 50 submarine (representing the `horizontal` offset, the `depth`, and also the 51 `aim` which is not used), return the product of the final value of the 52 `horizontal` offset and `depth`""" 53 return running_state_df.iloc[-1][['depth', 'horizontal']].prod() 54 55 def solve_puzzle(input_string): 56 """Return the numeric solution to the puzzle""" 57 # Wouldn't it be lovely if this instead looked like: 58 # return get_puzzle_input(2) |> 59 # convert_input_to_df() |> 60 # split_command_parts() |> 61 # convert_commands_to_offsets() |> 62 # find_submarine_state() |> 63 # calculate_puzzle_solution() 64 # Python needs a pipe operator! 65 return calculate_puzzle_solution( 66 find_submarine_state( 67 convert_commands_to_offsets( 68 split_command_parts( 69 convert_input_to_df(input_string) 70 ) 71 ) 72 ) 73 ) 74 75 if __name__ == "__main__": 76 print("Product of (new) horizontal and vertical offsets:", 77 solve_puzzle(get_puzzle_input(2)))