#!/usr/bin/env python3
import collections
import copy
SAMPLE = False
def load_input():
with open("sample.txt" if SAMPLE else "input.txt") as f:
return [line.strip() for line in f.readlines()]
def main():
displays = load_input()
result = 0
for i, display in enumerate(displays):
print(f"\nNew display #{i}: {display}")
all_digits, output_digits = display.split(" | ")
all_digits = [
''.join(sorted(digit))
for digit
in all_digits.split()
]
output_digits = [
''.join(sorted(digit))
for digit
in output_digits.split()
]
def digits_with_n_segments(n):
return [digit for digit in all_digits if len(digit) == n]
def digit_with_n_segments(n):
digits = digits_with_n_segments(n)
assert len(digits) == 1
return digits[0]
def digits_with_n_segments_containing_other_digit(n, other):
return [
digit
for digit in digits_with_n_segments(n)
if all([segment in digit for segment in other])
]
def digit_with_n_segments_containing_other_digit(n, other):
digits = digits_with_n_segments_containing_other_digit(n, other)
assert len(digits) == 1
return digits[0]
def digit_with_n_segments_containing_other_digit_but_not_this(n, other,
not_me):
digits = [
digit
for digit in digits_with_n_segments_containing_other_digit(n,
other)
if digit != not_me
]
assert len(digits) == 1
return digits[0]
def digits_with_n_segments_in_other_digit(n, other):
return [
digit
for digit in digits_with_n_segments(n)
if all([segment in other for segment in digit])
]
def digit_with_n_segments_in_other_digit_but_not_this(n, other, not_me):
digits = [
digit
for digit in digits_with_n_segments_in_other_digit(n, other)
if digit != not_me
]
assert len(digits) == 1
return digits[0]
def digit_with_n_segments_but_not_these(n, not_these):
digits = [
digit
for digit in digits_with_n_segments(n)
if digit not in not_these
]
assert len(digits) == 1
return digits[0]
one = digit_with_n_segments(2)
seven = digit_with_n_segments(3)
four = digit_with_n_segments(4)
eight = digit_with_n_segments(7)
three = digit_with_n_segments_containing_other_digit(5, one)
nine = digit_with_n_segments_containing_other_digit(6, three)
five = digit_with_n_segments_in_other_digit_but_not_this(5, nine, three)
six = digit_with_n_segments_containing_other_digit_but_not_this(6, five,
nine)
two = digit_with_n_segments_but_not_these(5, [three, five])
zero = digit_with_n_segments_but_not_these(6, [six, nine])
decoder = {
one: "1",
two: "2",
three: "3",
four: "4",
five: "5",
six: "6",
seven: "7",
eight: "8",
nine: "9",
zero: "0",
}
output_number = int(''.join(decoder[digit] for digit in output_digits))
assert output_number < 10000
result += output_number
print(result)
if __name__ == "__main__":
main()