day03_part2.py (2952B)
1 #!/usr/bin/env python 2 """Get the 'life support rating' by finding the 'oxygen generator rating' and 3 'CO2 scrubber tating' from the same stream of binary numbers used for part 4 1.""" 5 6 import functools 7 from utils import get_puzzle_input 8 from day03_part1 import (convert_input_to_df, 9 split_report_bits, 10 bit_series_to_int) 11 12 def filter_by_bit_commonality_at(bits_df, column_label, least_common=False): 13 """Helper function: return the rows for which the given `column_label` 14 has the most common bit (or least common). Ties go to the `1`s for most 15 common (`0` for least common)""" 16 # If we select all rows for which the value of the given column equals 17 # the integer value of the boolean corresponding to the test of whether 18 # twice the sum of the column is greater than or equal to the column 19 # length, we'll follow the rules for selecting the most common bit, and 20 # if we invert that, the least common 21 if len(bits_df) == 1: 22 return bits_df 23 rows_to_get = bits_df.loc[:, column_label] == \ 24 int(2 * bits_df.loc[:, column_label].sum() >= bits_df.loc[:, 1].size) 25 if least_common: 26 rows_to_get = ~rows_to_get 27 return bits_df.loc[rows_to_get] 28 29 def filter_by_bit_commonality(bits_df, least_common=False): 30 """Filter the df from left to right, selecting only rows that have the most 31 common bit (or least common, if `least_common` == `True`), until only a 32 single row remains, which is returned as a Series""" 33 # Using `reduce` instead of writing a loop prevents us from stopping as 34 # early as possible (with a break) and from checking that we have a unique 35 # solution, but it's just so CUTE I'm sorry 36 return functools.reduce( 37 lambda x, y: filter_by_bit_commonality_at(x, y, least_common), 38 bits_df.columns, 39 bits_df 40 ).iloc[0] 41 42 def calculate_oxygen_co2(bits_df): 43 """Given a df find the "oxygen generator rating" and "CO2 scrubber rating" 44 by successively filtering the df by the most (least) common bit from left 45 to right, and return the values as a tuple""" 46 return (bit_series_to_int(filter_by_bit_commonality(bits_df, False)), 47 bit_series_to_int(filter_by_bit_commonality(bits_df, True))) 48 49 def calculate_life_support_rating(oxygen_generator_rating, co2_scrubber_rating): 50 """Return the product of the oxygen generator rating and co2 scrubber 51 rating""" 52 return oxygen_generator_rating * co2_scrubber_rating 53 54 def solve_puzzle(input_string): 55 """Return the numeric solution to the puzzle""" 56 # 13.3 ms ± 131 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) 57 return calculate_life_support_rating( 58 *calculate_oxygen_co2( 59 split_report_bits( 60 convert_input_to_df(input_string) 61 ) 62 ) 63 ) 64 65 if __name__ == "__main__": 66 print("Life support rating:", 67 solve_puzzle(get_puzzle_input(3)))