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)))