advent_8.py

def a_minus_b(a, b):
  output = []
  for char in a:
    if char in b:
      continue
    output.append(char)
  return output

def solve(train_data):
  digits = train_data.split(' ')
  buckets = [None] * 8
  for digit in digits:
    if not buckets[len(digit)]:
      buckets[len(digit)] = []
    buckets[len(digit)].append(''.join(sorted(digit)))
  seven = buckets[3][0]
  one = buckets[2][0]
  segment_6 = a_minus_b(seven, one)[0]
  segment_one_or_two1, segment_one_or_two2 = one
  segment_1 = None
  segment_2 = None
  for digit in buckets[6]:
    if not segment_one_or_two1 in digit:
      segment_1 = segment_one_or_two1
      segment_2 = segment_one_or_two2
      break
    elif not segment_one_or_two2 in digit:
      segment_1 = segment_one_or_two2
      segment_2 = segment_one_or_two1
      break

  six = None
  for digit in buckets[6]:
    if not segment_1 in digit:
      six = digit
      break
  two_three_or_five1, two_three_or_five2, two_three_or_five3 = buckets[5]
  segment_3_or_7 = []
  for char in two_three_or_five1:
    if char == segment_6:
      continue
    if char in two_three_or_five2 and char in two_three_or_five3:
      segment_3_or_7.append(char)
  segment_3 = None
  segment_7 = None
  for digit in buckets[6]:
    if digit == six:
      continue
    if not segment_3_or_7[0] in digit:
      segment_7, segment_3 = segment_3_or_7
      break
    if not segment_3_or_7[1] in digit:
      segment_3, segment_7 = segment_3_or_7
      break

  four = buckets[4][0]
  segment_5 = None
  for char in four:
    if char in [segment_1, segment_2, segment_7]:
      continue
    segment_5 = char
  
  eight = buckets[7][0]
  segment_4 = None
  for char in eight:
    if char not in [segment_1, segment_2, segment_3, segment_5, segment_6, segment_7]:
      segment_4 = char
      break

  digit_map = {}
  for digit in buckets[6]:
    if segment_7 in digit:
      if segment_4 in digit:
        digit_map[digit] = 6
      else:
        digit_map[digit] = 9
    else:
      digit_map[digit] = 0
  for digit in buckets[5]:
    if segment_4 in digit:
      digit_map[digit] = 2
    elif segment_1 in digit:
      digit_map[digit] = 3
    else:
      digit_map[digit] = 5
  digit_map[one] = 1
  digit_map[four] = 4
  digit_map[seven] = 7
  digit_map[eight] = 8
  return digit_map

if __name__ == '__main__':
  with open('input.txt') as data:
    counts = [0] * 10
    total = 0
    for line in data:
      line = line.strip()
      train_data, digits = line.split(' | ')
      digit_map = solve(train_data)
      digits = digits.split(' ')
      line = 0
      for idx, digit in enumerate(digits):
        number = digit_map[''.join(sorted(digit))]
        counts[number] += 1
        line += number * (10 ** (len(digits) - 1 - idx))
      total += line
    print(counts[1] + counts[4] + counts[7] + counts[8], total)