AOC2020Day16inPython

import re
import math

# Part 0
# ------
# Parsing the data

# "131,103,...,197" -> (131,103,...,197)
def ticket_tuple(t): return tuple(map(int, t.split(",")))

# "price: 33-635 or 700-740" -> "price"
def field_name(f): return f.split(':')[0]
# "price: 33-635 or 700-740" -> set of valid values
def field_range(f):
    valid = set()
    for a,b in re.findall("(\d+)-(\d+)",f):
        valid |= set(range(int(a),int(b)+1))
    return valid

# Read in the data and parse it
# data[0] contains the fields
# data[1] contains my ticket (with a header row)
# data[2] contains all the other tickets (with a header row)

data =     [category.splitlines() for category in open("data16.txt").read().split("\n\n")]
myticket = ticket_tuple(data[1][1])
tickets  = [ticket_tuple(t) for t in data[2][1:]]
valid    = {field_name(f): field_range(f) for f in data[0]}

# Part 1
# ------
# Find the sum of all the invalid values
# Also filter the tickets, creating a valid_tickets set

invalid_sum, valid_tickets = 0, set()
valid_tickets.add(myticket)

for ticket in tickets:
    okay = True
    for val in ticket:
        if all(val not in valid[name] for name in valid):
            invalid_sum += val
            okay = False
    if okay: valid_tickets.add(ticket)

print("Part 1:", invalid_sum)

# Part 2
# ------
# Determine which category maps to which index in the tickets

# First we generate 'to_allocate', a dictionary of the fields we need to allocate.
# Each field is associated with a set of possible valid ticket indices

to_allocate = {f: set() for f in valid.keys()}
for i in range(len(myticket)):
    possibles = set(valid.keys())
    for t in valid_tickets:
        possibles = set(p for p in possibles if t[i] in valid[p])
    for p in possibles:
        to_allocate[p].add(i)

# We iterate through the list of unallocated fields, looking for ones which only
# have one possible index. Once we've found one, we allocate it, and remove it from
# the to_allocate dictionary, along with the index we've allocated.
# This continues until we've paired together all the fields and indices.

# This is 'greedy matching' and will only work if backtracking isn't required,
# in which case we'd need a more horrible algorithm - but luckily everything works!

allocated = {}
while to_allocate:
    for field in to_allocate:
        if len(to_allocate[field]) == 1:
            allocated[field] = to_allocate[field].pop()
            to_allocate.pop(field)
            for f in to_allocate: to_allocate[f].remove(allocated[field])
            break

# 'allocated' is now our mapping from field name to ticket index
# The actual calculation we need to do is to multiply together all the values
# in fields which have 'departure' in the field name.

print("Part 2:",math.prod(myticket[i] for (f,i) in allocated.items() if "departure" in f))