advent_of_code_2021

day05_part1.py (3743B)

---

 #!/usr/bin/env python
 """Advent of Code 2021, day 5 (part 1): the hydrothermal vents. Given the start
 and end coordinates of lines of vents, get information about the field."""
 
 # For Part 1, we don't need to operate on any 2D representation of the field
 # directly, so I won't bother building one. Again I'm sticking to pandas and
 # doing as much as I can with DFs.
 
 import numpy as np
 import pandas as pd
 from utils import get_puzzle_input
 
 def convert_input_to_df(input_string):
     """The input is a series of start and ending coordinates. Parse these and
     return a df with the columns `x1`, `y1`, `x2`, `y2`"""
     return (
         pd.Series([x for x in input_string.split('\n') if x])
         .str.split(',| -> ', n=4, expand=True)
         .rename(columns={0: 'x1', 1: 'y1', 2: 'x2', 3: 'y2'})
         .astype(np.int32)
     )
 
 def find_points_x(x_1, y_1, x_2, y_2):
     """Given a line segment's start and end coordinates, return a numby array
     of x coordinates for the points along segment"""
     # This is easy because we're only dealing with horizontal and vertical
     # lines (for now)
     if x_1 == x_2:
         return np.repeat(x_1, abs(y_2 - y_1) + 1)
     return np.arange(min(x_1, x_2), max(x_1, x_2)+1)
 
 def find_points_y(x_1, y_1, x_2, y_2):
     """Given a line segment's start and end coordinates, return a numby array
     of x coordinates for the points along segment"""
     return find_points_x(x_1=y_1, y_1=x_1, x_2=y_2, y_2=x_2)
 
 def find_line_points(endpoints_df):
     """Given a df with the start and end coordinates of the line segments, add
     new columns, `x` and `y`, that give the coordinates of all points covered
     by the line (note that line data will be dropped but indices are
     retained)"""
     # In Part 1 (I sense that Part 2 will be different), we only transform the
     # horizontal and verticl line segments. We'll remove the diagonal line
     # segments from the df.
     return (
         endpoints_df
         .loc[(endpoints_df['x1'] == endpoints_df['x2']) |
              (endpoints_df['y1'] == endpoints_df['y2'])]
         .assign(x=lambda df: [find_points_x(*x[1].to_list()) \
                               for x in df[['x1', 'y1', 'x2', 'y2']].iterrows()],
                 y=lambda df: [find_points_y(*x[1].to_list()) \
                               for x in df[['x1', 'y1', 'x2', 'y2']].iterrows()])
         .explode(['x', 'y'])
         .drop(columns=['x1', 'y1', 'x2', 'y2'])
     )
 
 def count_point_coverage(points_df):
     """Given a df with the `x` and `y` coordinates of all covered points, find
     out how many times each point is covered (among points covered at least
     once)"""
     return (
         points_df
         .groupby(['x', 'y']).size()
         .reset_index()
         .rename(columns={0: 'n'})
     )
 
 def count_danger_points(coverage_df):
     """Given a df with the total coverage for each coordinate, count the number
     of points with a coverage > 2"""
     return len(coverage_df.loc[coverage_df['n'] > 1])
 
 def solve_puzzle(input_string):
     """Return the numeric solution to the puzzle"""
     return count_danger_points(
         count_point_coverage(
             find_line_points(
                 convert_input_to_df(input_string)
             )
         )
     )
 
 if __name__ == "__main__":
     assert solve_puzzle(
         "\n".join(("0,9 -> 5,9",
                    "8,0 -> 0,8",
                    "9,4 -> 3,4",
                    "2,2 -> 2,1",
                    "7,0 -> 7,4",
                    "6,4 -> 2,0",
                    "0,9 -> 2,9",
                    "3,4 -> 1,4",
                    "0,0 -> 8,8",
                    "5,5 -> 8,2",
                    "\n"))) == 5
 
     print("Number of points to avoid:",
           solve_puzzle(get_puzzle_input(5)))