aoc/2022/16/solution.py

#!/usr/bin/python3
import sys
import re
import functools

# part 1

valves = []
rates = []
indices = []
adjacent = []
index_map = {}

def parse_input():
    nodes = 0
    # process input
    for line in sys.stdin:
        source, rate, targets = re.match(r'^Valve ([A-Z]+) has flow rate=([0-9]+); tunnels? leads? to valves? ([A-Z, ]+)', line.strip()).groups()
        index_map[source] = nodes
        valves.append(nodes)
        rates.append(int(rate))
        adjacent.append(targets.split(', '))
        nodes += 1

    # transform letter neighbors to index neighbors
    for i in range(nodes):
        indices.append([*map(lambda x: index_map[x], adjacent[i])])
    #print(f'index_map {index_map}\nadjacent {adjacent}')
    print(f'nodes {nodes}\nvalves {valves}\nrates {rates}\nindices {indices}')
    #return nodes, index_map, valves, rates, indices

#@functools.lru_cache(maxsize=None)
def enumerate_pressure(seen, current, time=30):
    print('looking at', current)
    if time <= 0:
        return 0
    
    max_pressure = 0
    if current not in seen:
        pressure = rates[current] * (time - 1)
        new_seen = seen.union({current})
        print(f'{current} not yet seen, pressure {pressure} (rate {rates[current]}, time {time - 1})')
        for index in indices[current]:
            if pressure == 0:
                max_pressure = max(max_pressure, enumerate_pressure(new_seen, index, time - 1))
            else:
                max_pressure = max(max_pressure, pressure + enumerate_pressure(new_seen, index, time - 2))

    #print(seen, current, time, max_pressure)
    return max_pressure

def part1():
    # parse input
    parse_input()
    # find maximal pressure
    #for i in range(nodes):
        #visited = [False for i in range(nodes)]
        #print(maximal_pressure(valves, rates, indices, visited, i))
    seen = frozenset()
    print(enumerate_pressure(seen, 0))

# part 2

def part2():
    pass

if sys.argv[1] in '1':
    part1()
else:
    part2()