advent_of_code_2021

day08_part2.py (5026B)

---

 #!/usr/bin/env python
 """Advent of Code 2021, day 8 (part 2): decode all of the scrambled
 seven-segment displays"""
 
 # This is a fun puzzle. I sat down with pen and paper and proved to myself that
 # you can deduce the mapping from scrambled segments to the full set of digits
 # using just the information provided. It looks like this:
 # * If 2 segments are lit, that's '1'
 # * If 3 segments are lit, that's '7'
 # * If 4 segments are lit, that's '4'
 # * If 7 segments are lit, that's '8'
 # (So far, this was given in the puzzle prompt, nothing new yet)
 # * If 5 segments are lit, it's '2', '3', or '5'
 # * If 6 segments are lit, it's '0', '6', or '9'
 # (Okay, that's not useful... yet)
 # * If 5 segments are lit and they include all the segments in '7', that's '3'
 # * If 6 segments are lit and they don't include all the segments in '7',
 #   that's '6'
 # * If 6 segments are lit and they include the union of the segments comprising
 #   '3' and '4', that's '9'
 # * If 6 segments are lit and it's not a '6' or '9', that's '0'
 # * If 5 segments are lit and it is missing only one of the segments from '6',
 #   that's '5'
 # * If 5 segments are lit and it's not a '2' or a '3', that's '2'
 # It's part 2, so apparently we're doing it!
 
 from day08_part1 import (EXAMPLE_INPUT,
                          convert_input_to_df,
                          apply_mappers_to_digits)
 from utils import get_puzzle_input
 
 def create_mapper_from_segments(segment_sequence):
     """Given the observed sequence of segment sets (which isn't actually used
     for anything here), return a function that maps a `frozenset` of segments
     to a digit"""
 
     # This approach isn't efficient (or elegant) at all, but I don't think this
     # function is being run enough times, on a long enough list, to justify the
     # code complexity for a more efficient/elegant approach
     segment_sets = [frozenset(i) for i in segment_sequence.split(' ')]
     digits_to_segments = {}
 
     # Step 1 of building a segment to digit map: find the mapping to '1', '7',
     # '4', and '8'
     for i, seg in enumerate(segment_sets):
         if len(seg) == 2:
             digits_to_segments['1'] = segment_sets.pop(i)
             break
     for i, seg in enumerate(segment_sets):
         if len(seg) == 3:
             digits_to_segments['7'] = segment_sets.pop(i)
             break
     for i, seg in enumerate(segment_sets):
         if len(seg) == 4:
             digits_to_segments['4'] = segment_sets.pop(i)
             break
     for i, seg in enumerate(segment_sets):
         if len(seg) == 7:
             digits_to_segments['8'] = segment_sets.pop(i)
             break
 
     # Step 2: using '7' alone, find '3' and '6'
     for i, seg in enumerate(segment_sets):
         if len(seg) == 5 and not digits_to_segments['7'] - seg:
             digits_to_segments['3'] = segment_sets.pop(i)
             break
     for i, seg in enumerate(segment_sets):
         if len(seg) == 6 and (digits_to_segments['7'] - seg):
             digits_to_segments['6'] = segment_sets.pop(i)
             break
 
     # Step 3: Using '4', '3', and '6', find '9', and '5'
     for i, seg in enumerate(segment_sets):
         if len(seg) == 5 and len(digits_to_segments['6'] - seg) == 1:
             digits_to_segments['5'] = segment_sets.pop(i)
             break
     for i, seg in enumerate(segment_sets):
         if len(seg) == 6 and \
                 seg == digits_to_segments['3'] | digits_to_segments['4']:
             digits_to_segments['9'] = segment_sets.pop(i)
             break
 
     # Step 4 (final): Using '3', '6', '2', and '9', find '0' and '2'
     for i, seg in enumerate(segment_sets):
         if len(seg) == 5:
             digits_to_segments['2'] = segment_sets.pop(i)
             break
     digits_to_segments['0'] = segment_sets.pop()
 
     segments_to_digits = {v:k for k, v in digits_to_segments.items()}
 
     def part_2_mapper(segments):
         """Just a closure wrapping the dictionary (hackily) built above"""
         return segments_to_digits.get(segments)
     return part_2_mapper
 
 def convert_single_digits_to_integers(digits_df):
     """Given a four-column df of digit displays, return a Series consisting of
     a the single integer value obtained across each row"""
     # Can this be done with `apply` somehow? This feels needlessly inelegant
     return (
         digits_df['digit_1']
         .add(digits_df['digit_2'])
         .add(digits_df['digit_3'])
         .add(digits_df['digit_4'])
         .astype(int)
     )
 
 def solve_puzzle(input_string):
     """Return the numeric solution to the puzzle"""
     return int(
         convert_single_digits_to_integers(
             apply_mappers_to_digits(
                 convert_input_to_df(input_string),
                 create_mapper_from_segments
             )
         )
         .sum()
     )
 
 def main():
     """Run when the file is called as a script"""
     assert solve_puzzle(EXAMPLE_INPUT) == 61229
     print("Sum of all decoded displays:",
           solve_puzzle(get_puzzle_input(8)))
 
 if __name__ == "__main__":
     main()