day08_2.py

def contains_all(str, set):
    return 0 not in [c in str for c in set]


def find_digits(input_parts):
    digits = [""] * 10
    cands960 = set()
    cands532 = set()
    for input_part in input_parts:
        if len(input_part) == 2:
            digits[1] = input_part
        elif len(input_part) == 4:
            digits[4] = input_part
        elif len(input_part) == 3:
            digits[7] = input_part
        elif len(input_part) == 7:
            digits[8] = input_part
        elif len(input_part) == 6:
            cands960.add(input_part)
        elif len(input_part) == 5:
            cands532.add(input_part)
    digits[9] = next(filter(lambda x: contains_all(x, digits[4]), cands960))
    digits[3] = next(filter(lambda x: contains_all(x, digits[1]), cands532))
    digits[5] = next(filter(lambda x: contains_all(digits[9], x), cands532.difference({digits[3]})))
    digits[6] = next(filter(lambda x: contains_all(x, digits[5]), cands960.difference({digits[9]})))
    digits[2] = cands532.difference({digits[5], digits[3]}).pop()
    digits[0] = cands960.difference({digits[9], digits[6]}).pop()
    return list(map(lambda x: "".join(sorted(x)), digits))


def calculate_line_value(line):
    parts = line.split("|")
    input_parts = parts[0].split()
    output_parts = parts[1].split()
    digits = find_digits(input_parts)
    digit_values = [digits.index("".join(sorted(op))) for op in output_parts]
    return sum([digit_values[i] * pow(10, len(digit_values) - i - 1) for i in range(0, len(digit_values))])


if __name__ == '__main__':
    with open("input.txt") as file:
        lines = [line.rstrip() for line in file.readlines()]
    tot = sum([calculate_line_value(line) for line in lines])
    assert tot == 968175