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

day05_part2.py (2715B)


      1 #!/usr/bin/env python
      2 """Advent of Code 2021, day 5 (part 2): considering the diagonal lines."""
      3 
      4 import numpy as np
      5 from utils import get_puzzle_input
      6 from day05_part1 import (convert_input_to_df,
      7                          count_point_coverage,
      8                          count_danger_points)
      9 
     10 def find_points_x(x_1, y_1, x_2, y_2):
     11     """Given a line segment's start and end coordinates, return a numby array
     12     of x coordinates for the points along segment"""
     13     # This is easy because we're only dealing with horizontal and vertical
     14     # lines (for now)
     15     if x_1 > x_2:
     16         line_array = np.arange(x_1, x_2 - 1, -1)
     17     elif x_1 < x_2:
     18         line_array = np.arange(x_1, x_2 + 1, 1)
     19     else:
     20         line_array = np.repeat(x_1, abs(y_2 - y_1) + 1)
     21     return line_array
     22 
     23 def find_points_y(x_1, y_1, x_2, y_2):
     24     """Given a line segment's start and end coordinates, return a numby array
     25     of x coordinates for the points along segment"""
     26     return find_points_x(x_1=y_1, y_1=x_1, x_2=y_2, y_2=x_2)
     27 
     28 def find_line_points(endpoints_df):
     29     """Given a df with the start and end coordinates of the line segments, add
     30     new columns, `x` and `y`, that give the coordinates of all points covered
     31     by the line (note that line data will be dropped but indices are
     32     retained)"""
     33     # I don't love that I'm copy-pasting so much of the code from Part 1, but
     34     # one thing I have consciously avoided is refactoring Part 1 after Part 2
     35     # is revealed to me. Not sure why I care about that but I do!
     36     return (
     37         endpoints_df
     38         .assign(x=lambda df: [find_points_x(*x[1].to_list()) \
     39                               for x in df[['x1', 'y1', 'x2', 'y2']].iterrows()],
     40                 y=lambda df: [find_points_y(*x[1].to_list()) \
     41                               for x in df[['x1', 'y1', 'x2', 'y2']].iterrows()])
     42         .explode(['x', 'y'])
     43         .drop(columns=['x1', 'y1', 'x2', 'y2'])
     44     )
     45 
     46 def solve_puzzle(input_string):
     47     """Return the numeric solution to the puzzle"""
     48     return count_danger_points(
     49         count_point_coverage(
     50             find_line_points(
     51                 convert_input_to_df(input_string)
     52             )
     53         )
     54     )
     55 
     56 if __name__ == "__main__":
     57     assert solve_puzzle(
     58         "\n".join(("0,9 -> 5,9",
     59                    "8,0 -> 0,8",
     60                    "9,4 -> 3,4",
     61                    "2,2 -> 2,1",
     62                    "7,0 -> 7,4",
     63                    "6,4 -> 2,0",
     64                    "0,9 -> 2,9",
     65                    "3,4 -> 1,4",
     66                    "0,0 -> 8,8",
     67                    "5,5 -> 8,2",
     68                    "\n"))) == 12
     69 
     70     print("Number of points to avoid (including the diagnoal lines now):",
     71           solve_puzzle(get_puzzle_input(5)))