day02_part1.py (3015B)
1 #!/usr/bin/env python 2 """Day 2, Part 1: we need to parse submarine commands and figure out where the 3 submarine winds up. I am very confident that I wrote code to solve this exact 4 problem, using Lisp, decades ago in undergrad.""" 5 6 # Once again I'm going to take a pandas-heavy approach even though it's not 7 # strictly necessary here. I'm going to make some tweaks to the coding 8 # conventions I used yesterday; I'd like to make it easier to develop different 9 # implementations and compare their performance. 10 11 from io import StringIO 12 import numpy as np 13 import pandas as pd 14 from utils import get_puzzle_input 15 16 def convert_input_to_df(input_string): 17 """Given a string representation of the input data, return a single-column 18 pandas data frame with the column `commands`""" 19 return pd.read_csv(StringIO(input_string), 20 names=('commands',)) 21 22 def split_command_parts(command_df): 23 """Given a pandas data frame with a column `commands`, returns a new data 24 frame with the columns `direction` and `distance`""" 25 return ( 26 command_df['commands'] 27 .str.split(n=2, expand=True) 28 .rename(columns={0:'direction', 1:'distance'}) 29 .assign(distance=lambda x: pd.to_numeric(x['distance'])) 30 ) 31 32 def convert_commands_to_offsets(command_df): 33 """Given a pandas data frame with the columns `direction` (in "forward", 34 "up", and "down") and `distance`, returns a new data frame with the columns 35 `axis` (either "horizontal" or "vertical") and `offset`""" 36 return pd.DataFrame( 37 {'axis': np.where(command_df['direction'].isin(('up', 'down')), 38 'vertical', 'horizontal'), 39 'offset': np.where(command_df['direction'] == 'up', 40 -command_df['distance'], command_df['distance'])} 41 ) 42 43 def find_total_offsets(offset_df): 44 """Given a pandas data frame with the column `axis` and any others, find 45 the sum of other columns grouped by `axis`. (Note that this doesn't 46 restrict itself to the axes 'horizontal' or 'vertical' and it doesn't check 47 the other column names""" 48 # You know, in retrospect it was maybe confusing to call a column 'axis' in 49 # pandas world. Ah well! 50 return offset_df.groupby('axis').sum() 51 52 def calculate_puzzle_solution(axis_summary_df): 53 """Given a pandas data frame with the sum of `offset` values along each 54 `axis`, multiply the 'horizontal' and 'vertical' offsets.""" 55 return axis_summary_df.loc['horizontal', 'offset'] * \ 56 axis_summary_df.loc['vertical', 'offset'] 57 58 def solve_puzzle(input_string): 59 """Return the numeric solution to the puzzle""" 60 return calculate_puzzle_solution( 61 find_total_offsets( 62 convert_commands_to_offsets( 63 split_command_parts( 64 convert_input_to_df(input_string) 65 ) 66 ) 67 ) 68 ) 69 70 if __name__ == "__main__": 71 print("Product of horizontal and vertical offsets:", 72 solve_puzzle(get_puzzle_input(2)))