#!/usr/bin/env python3 import sys, re from copy import deepcopy VALVES = {'V': [], 'E': [], 'w': {}} def parse(line): g = re.compile(r"Valve (?P\w{2}) has flow rate=(?P\d+); tunnel(s)? lead(s)? to valve(s)? (?P(\w{2}(, )?)*)").search(line) vc = g.group('valve') + 'C' vo = g.group('valve') + 'O' es = g.group('tunnel').split(', ') VALVES['V'] += [vc,vo] VALVES['E'].append((vc,vo)) for e in es: VALVES['E'].append((vo,e+'C')) VALVES['E'].append((vc,e+'C')) VALVES['w'][(vo,e+'C')] = 0 VALVES['w'][(vc,e+'C')] = 0 VALVES['w'][(vc,vo)] = int(g.group('fr')) def Init(G,s,k): d = {} o = {} for v in G['V']: d[v] = [-1] * k o[v] = [[]] * k d[s] = [0] * k return d, o def anti_relax(u,v,w,i,d,k,o): if d[v][i] < d[u][i-1] + w[(u,v)] * (k-i): d[v][i] = d[u][i-1] + w[(u,v)] * (k-i) o[v][i] = o[u][i-1] + [u] def find_path(G,s,k): d,o = Init(G,s,k) for i in range(k): di = deepcopy(d) for (u,v) in G['E']: if d[u][i-1] > -1 and not (u[:2] == v[:2] and v in o[u][i-1]): anti_relax(u,v,G['w'],i,di,k,o) d = deepcopy(di) return max(map(max,d.values())) def Init2(G,s,k): d = {} o = {} for v in G['V']: d[v] = [-1] * k o[v] = [[]] * k d[s] = [0] * k return d, o def anti_relax2(u,v,w,i,d,k,o): if d[v][i] < d[u][i-1] + w[(u,v)] * (k-i): d[v][i] = d[u][i-1] + w[(u,v)] * (k-i) o[v][i] = o[u][i-1] + eval(u) def find_path2(G,s,k): d,o = Init2(G,s,k) for i in range(k): di = deepcopy(d) for (u,v) in G['E']: u1,u2 = eval(u) v1,v2 = eval(v) prob = False for (a,b) in [(u1,v1),(u1,v2),(u2,v1),(u2,v2)]: if a[:2] == b[:2]: if a in o[u][i-1]: prob = True if d[u][i-1] > -1 and not prob: anti_relax2(u,v,G['w'],i,di,k,o) d = deepcopy(di) return max(map(max,d.values())) def find_double_path(G,s,k): G2 = {'V': [], 'E': [], 'w': {}} for u in G['V']: for v in G['V']: if u == v and u[2] == 'O': continue G2['V'].append(str(sorted([u,v]))) for e in G['E']: if e[0] == u: for f in G['E']: if e[1] == f[1] and e[1][2] == 'O': continue if f[0] == v: G2['E'].append((str(sorted([u,v])),str(sorted([e[1],f[1]])))) G2['w'][(str(sorted([u,v])),str(sorted([e[1],f[1]])))] = G['w'][e] + G['w'][f] return find_path2(G2,f"['{s}', '{s}']",k) if __name__ == '__main__': [parse(line.strip('\n')) for line in open(sys.argv[1])] m = find_path(VALVES,'AAC',29) # start (s): AA, length (k): 30, # challenge 1 res1 = str(m) print("challenge 1:" + "\n" + res1 + "\n") # challenge 2 m = find_double_path(VALVES,'AAC',25) # missing: magic! Doesn't work for example! (off by one) But does work for input... res2 = str(m) print("challenge 2:" + "\n" + res2 + "\n")