commit 235c28b62909ff016963880a7034ac465ee440df
parent 49b32ed7252e5edd87e0a4b426787193bcd5442f
Author: Eamon Caddigan <eamon.caddigan@gmail.com>
Date: Thu, 2 Dec 2021 14:09:09 -0500
Solution to day 2, part 2
Diffstat:
| A | day02_part2.py | | | 69 | +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ |
1 file changed, 69 insertions(+), 0 deletions(-)
diff --git a/day02_part2.py b/day02_part2.py
@@ -0,0 +1,69 @@
+#!/usr/bin/env python
+"""Day 2, Part 2: It turns out that controlling the submarine is more
+complicated than we thought in Part 1. We need to maintain a third state
+variable (in addition to horizontal and vertical position) called 'aim', and
+use that when calculating sub movement."""
+
+# Hey, the functions from Part 1 will be useful once again. All the input
+# parsing is the same, and now we just need to use slightly more complicated
+# logic when figuring out final position. We can start from the point where we
+# have a data frame with the columns `axis` and `offset`.
+
+#from io import StringIO
+import numpy as np
+import pandas as pd
+from utils import get_puzzle_input
+from day02_part1 import (convert_input_to_df,
+ split_command_parts,
+ convert_commands_to_offsets)
+
+def find_submarine_state(offset_df):
+ """Given a pandas data frame with the columns `axis` and `offset`, we track
+ `aim` as well as `horizontal` and `vertical` offsets, using complicated
+ rules:
+ * 'vertical X' increases or decreases aim (which starts at 0)
+ * 'horizontal X':
+ * increases horizontal position by X
+ * increases vertical position (i.e., depth) by aim * X
+ This returns a data frame with the columns `horizontal`, `vertical` and
+ `aim`, representing the running value of these states."""
+ return (
+ pd.DataFrame(
+ # 'aim' represents the up-to-date state of the aim parameter, while
+ # 'horizontal_step' is just the current step's offset. We save this
+ # intermediate calculation because we need it for 'depth'.
+ # This line uses the extremely new-to-me 'walrus operator' to
+ # cretae the variable `vert_rows` so we don't calculate the
+ # equality twice. Not sure how I feel about this style.
+ {'aim': np.where(vert_rows := offset_df['axis'] == 'vertical',
+ offset_df['offset'], 0).cumsum(),
+ 'horizontal_step': np.where(~vert_rows,
+ offset_df['offset'], 0)}
+ )
+ .assign(depth=lambda x: (x['aim'] * x['horizontal_step']).cumsum(),
+ horizontal=lambda x: x['horizontal_step'].cumsum())
+ .drop('horizontal_step', axis='columns')
+ )
+
+def calculate_puzzle_solution(running_state_df):
+ """Given a pandas data frame with the running values of the state of the
+ submarine (representing the `horizontal` offset, the `depth`, and also the
+ `aim` which is not used), return the product of the final value of the
+ `horizontal` offset and `depth`"""
+ return running_state_df.iloc[-1][['depth', 'horizontal']].prod()
+
+def solve_puzzle():
+ """Return the numeric solution to the puzzle"""
+ return calculate_puzzle_solution(
+ find_submarine_state(
+ convert_commands_to_offsets(
+ split_command_parts(
+ convert_input_to_df(get_puzzle_input(2))
+ )
+ )
+ )
+ )
+
+if __name__ == "__main__":
+ print("Product of (new) horizontal and vertical offsets:",
+ solve_puzzle())
