vrazix

Untitled

Dec 8th, 2021
1,103
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. from collections import Counter, defaultdict
  2. from enum import IntEnum
  3.  
  4. raw = [line.split(' | ') for line in f_input.split('\n')]
  5.  
  6. all_unique, outputs = zip(*raw)
  7.  
  8. all_unique = [un.split() for un in all_unique]
  9. outputs = [output.split() for output in outputs]
  10. flat_outputs = [item for sublist in outputs for item in sublist]
  11.  
  12. output_counter = Counter(len(fo) for fo in flat_outputs)
  13.  
  14. print(sum([count for segments, count in output_counter.items() if segments in (2, 3, 4, 7)]))
  15.  
  16. # part 2
  17.  
  18. class D(IntEnum):
  19.     '''
  20.     000
  21.    1   2
  22.    1   2
  23.     333
  24.    4   5
  25.    4   5
  26.     666
  27.    '''
  28.     TOP = 0
  29.     TOPLEFT = 1
  30.     TOPRIGHT = 2
  31.     MIDDLE = 3
  32.     BOTTOMLEFT = 4
  33.     BOTTOMRIGHT = 5
  34.     BOTTOM = 6
  35.  
  36. ENCODER = {(D.TOP, D.TOPLEFT, D.TOPRIGHT, D.BOTTOMLEFT, D.BOTTOMRIGHT, D.BOTTOM): '0',
  37.            (D.TOPRIGHT, D.BOTTOMRIGHT): '1',
  38.            (D.TOP, D.TOPRIGHT, D.MIDDLE, D.BOTTOMLEFT, D.BOTTOM): '2',
  39.            (D.TOP, D.TOPRIGHT, D.MIDDLE, D.BOTTOMRIGHT, D.BOTTOM): '3',
  40.            (D.TOPLEFT, D.TOPRIGHT, D.MIDDLE, D.BOTTOMRIGHT): '4',
  41.            (D.TOP, D.TOPLEFT, D.MIDDLE, D.BOTTOMRIGHT, D.BOTTOM): '5',
  42.            (D.TOP, D.TOPLEFT, D.MIDDLE, D.BOTTOMLEFT, D.BOTTOMRIGHT, D.BOTTOM): '6',
  43.            (D.TOP, D.TOPRIGHT, D.BOTTOMRIGHT): '7',
  44.            (D.TOP, D.TOPLEFT, D.TOPRIGHT, D.MIDDLE, D.BOTTOMLEFT, D.BOTTOMRIGHT, D.BOTTOM): '8',
  45.            (D.TOP, D.TOPLEFT, D.TOPRIGHT, D.MIDDLE, D.BOTTOMRIGHT, D.BOTTOM): '9',
  46.            }
  47.  
  48. def encode_segments(segments, mapping):
  49.     '''mapping = {'a': D.TOPRIGHT, 'b': D.BOTTOMRIGHT, 'c': D.BOTTOM, 'd': D.TOP,
  50.                  'e': D.TOPLEFT, 'f': D.MIDDLE, 'g': D.BOTTOMLEFT}
  51.    encode_segments('acedgfb', mapping) => '8'
  52.    '''
  53.  
  54.     encodable = tuple(sorted(mapping[letter] for letter in segments))
  55.  
  56.     return ENCODER[encodable]
  57.  
  58.  
  59. def decode_segments(unique_specification):
  60.     '''['acedgfb', 'cdfbe', 'gcdfa', 'fbcad', 'dab', 'cefabd', 'cdfgeb', 'eafb', 'cagedb', 'ab']
  61.    => {'a': D.TOPRIGHT, 'b': D.BOTTOMRIGHT, 'c': D.BOTTOM, 'd': D.TOP,
  62.        'e': D.TOPLEFT, 'f': D.MIDDLE, 'g': D.BOTTOMLEFT}'''
  63.  
  64.     def create_length_dictionary(unique_specification):
  65.         '''u_s = ['acedgfb', 'cdfbe', 'gcdfa', 'fbcad', 'dab', 'cefabd', 'cdfgeb', 'eafb', 'cagedb', 'ab']
  66.        create_length_dictionary(u_s) => {2: ['ab'], 3: ['dab']: 4: ['eafb'], 5: ['cdfbe', 'gcdfa', 'fbcad],
  67.                                          6: ['cefabd', 'cdfgeb', 'cagedb'], 7: ['acedgfb']}
  68.        '''
  69.  
  70.         ldict = defaultdict(list)
  71.  
  72.         for letters in unique_specification:
  73.             ldict[len(letters)].append(letters)
  74.  
  75.         return ldict
  76.  
  77.     mapping = dict()
  78.  
  79.     ldict = create_length_dictionary(unique_specification)
  80.  
  81.     # we know these
  82.     the_one, the_seven = set(*ldict[2]), set(*ldict[3])
  83.     the_four, the_eight =  set(*ldict[4]), set(*ldict[7])
  84.  
  85.     # 3 segments (the seven) minus 2 (the one) segments yields D.TOP segment
  86.     (segment,) = the_seven - the_one
  87.     mapping[segment] = D.TOP
  88.  
  89.     # 4 segments (the four) minus 3 (the seven) segments yields D.TOPLEFT and D.MIDDLE
  90.     TL_M = the_four - the_seven
  91.  
  92.     # 7 segments (the eight) minus each of the three 6 segment configurations yields 1 segment each
  93.     three_1_seg = [the_eight - set(six) for six in ldict[6]]
  94.  
  95.     assert all(len(seg) == 1 for seg in three_1_seg)
  96.  
  97.     # we can decouple D.TOPLEFT and D.MIDDLE by checking each of the leftovers
  98.     # the one with an overlap is D.MIDDLE, the other is D.TOPLEFT
  99.     for segment in three_1_seg:
  100.  
  101.         if TL_M & segment:
  102.             (tl_seg,) = TL_M - segment
  103.             mapping[tl_seg] = D.TOPLEFT
  104.  
  105.             (segment,) = segment
  106.             mapping[segment] = D.MIDDLE
  107.  
  108.     # we now know D.TOP, D.TOPLEFT amd D.MIDDLE
  109.     # from the 5 segment results (digits 2, 3 and 5) and the digit 1, we can unlock much
  110.     # first, the one with BOTH segments in digit 1 is the three, and the unknown
  111.     # segment is D.BOTTOM
  112.  
  113.     for segments in ldict[5]:
  114.         r_unk = set(segments) - the_one
  115.  
  116.         if len(r_unk) == 3:
  117.             (segment,) = r_unk - set(mapping.keys())
  118.             mapping[segment] = D.BOTTOM
  119.             #the_three = segments
  120.  
  121.     # now if we subtract our knowns (D.TOP, D.TOPLEFT, D.MIDDLE, D.BOTTOM) from the 5 segment results, the one that isn't
  122.     # two results will yield the bottom right
  123.    
  124.     knowns = set(mapping.keys())
  125.     for segments in ldict[5]:
  126.         remain = set(segments) - knowns
  127.  
  128.         if len(remain) == 1:
  129.             (segment,) = remain
  130.             mapping[segment] = D.BOTTOMRIGHT
  131.  
  132.     # two left, D.TOPRIGHT and D.BOTTOMLEFT. top right is easy, knowns shared with the one
  133.  
  134.     (segment,) = the_one - set(mapping.keys())
  135.     mapping[segment] = D.TOPRIGHT
  136.  
  137.     (segment,) = set('abcdefg') - set(mapping.keys())
  138.     mapping[segment] = D.BOTTOMLEFT
  139.  
  140.     assert len(mapping) == 7
  141.  
  142.     return mapping
  143.  
  144. # okay, let's do this.
  145.  
  146. final_answer = 0
  147.  
  148. for all_u, output in zip(all_unique, outputs):
  149.  
  150.     mapping = decode_segments(all_u)
  151.     final_answer += int(''.join([encode_segments(sequence, mapping) for sequence in output]))
  152.  
  153. print(final_answer)
RAW Paste Data