adventofcode2022/day16/sol.py

G = []
for line in open("input"):
    words = line.split()
    valve = words[1]
    rate = words[4]
    edges = [x.strip(", ") for x in words[9:]]

    print(rate)
    rate = rate.strip(";")
    rate = rate[rate.index("=")+1:]
    rate = int(rate)
    G.append((valve, rate, edges))

print(G)

def save():
    seen = set()
    with open("input.dot", "w") as f:
        print("graph {", file=f)
        for v, r, E in G:
            if r:
                print('%s [label="%s (%s)"]' % (v, v, r), file=f)
            for e in E:
                if (e,v) not in seen:
                    print("%s -- %s" % (v,e), file=f)
                    seen.add((v,e))
        print("}", file=f)

#save()

#import astar

def part1():
    score = {'AA': (0, 0, [])}

    minutes = 30
    for _ in range(minutes):
        minutes -= 1
        next = {}
        for v, r, E in G:
            vo = v + 'o'
            s = []
            o = []
            if vo in score:
                reward, flow, open = score[vo]
                o.append((reward, flow, open))
            if v in score:
                reward, flow, open = score[v]
                # stay, don't open valve
                s.append((reward, flow, open))
                # stay in place, open valve
                if v not in open:
                    reward += r*minutes
                    o.append((reward, flow+r, open+[v]))
            # move here from somewhere else
            for e in E:
                if e in score:
                    if v in score[e][-1]:
                        o.append(score[e])
                    else:
                        s.append(score[e])
                eo = e+'o'
                if eo in score:
                    if v in score[eo][-1]:
                        o.append(score[eo])
                    else:
                        s.append(score[eo])
            if s:
                next[v] = max(s)
            if o:
                next[vo] = max(o)
        score = next
        print("%d minutes left" % minutes)
        for v, (r, flow, open) in sorted(score.items(), key=lambda x: x[1]):
            print("\t", v, r, "\t", ",".join(open))
    print(max(r for r,_,_ in score.values()))

def part2():
    H = {x[0]: x for x in G}
    def actions(v, open):
        _, r, E = H[v]
        if v in open:
            yield (v, 0, [])
        else:
            yield (v, 0, [])
            yield (v+'o', r, [v])
        # move to somewhere else
        for e in E:
            if e in open:
                e = e+'o'
            yield (e, 0, [])

    score = {('AA', 'AA'): (0, 0, [])}
    minutes = 26
    for _ in range(minutes):
        next = {}
        minutes -= 1

        for (me, elephant) in score:
            reward, flow, open = score[(me, elephant)]
            for v, r, o in actions(me.strip('o'), open):
                for v1, r1, o1 in actions(elephant.strip('o'), open+o):
                    reward_ = reward + (r+r1)*minutes
                    flow_ = flow + r + r1
                    open_ = open + o + o1
                next.setdefault((v, v1), []).append((reward_, flow_, open_))

        for k in next:
            next[k] = max(next[k])

        score = next
        print("%d minutes left" % minutes)
        for v, (r, flow, open) in sorted(score.items(), key=lambda x: x[1]):
            print("\t", v, r, "\t", ",".join(open))
    print(max(r for r,_,_ in score.values()))

#part2()

def part2_b():
    score = {'AA': (0, 0, [])}

    minutes = 26
    for _ in range(minutes):
        minutes -= 1
        next = {}
        for v, r, E in G:
            vo = v + 'o'
            s = []
            o = []
            if vo in score:
                reward, flow, open = score[vo]
                o.append((reward, flow, open))
            if v in score:
                reward, flow, open = score[v]
                # stay, don't open valve
                s.append((reward, flow, open))
                # stay in place, open valve
                if v not in open:
                    reward += r*minutes
                    o.append((reward, flow+r, open+[v]))
            # move here from somewhere else
            for e in E:
                if e in score:
                    if v in score[e][-1]:
                        o.append(score[e])
                    else:
                        s.append(score[e])
                eo = e+'o'
                if eo in score:
                    if v in score[eo][-1]:
                        o.append(score[eo])
                    else:
                        s.append(score[eo])
            if s:
                next[v] = max(s)
            if o:
                next[vo] = max(o)
        score = next
        print("%d minutes left" % minutes)
        for v, (r, flow, open) in sorted(score.items(), key=lambda x: x[1]):
            print("\t", v, r, "\t", ",".join(open))

    maxpair = []
    for r,_,open in score.values():
        o = frozenset(open)
        for s,_,open in score.values():
            if o.isdisjoint(open):
                maxpair.append(r+s)
    print(max(maxpair))

#part2_b()

def part2_c():
    V = sorted(v for v,_,_ in G)
    B = {v: 1<<i for i, v in enumerate(V)}
    E = {B[v]: [B[e] for e in edges] for v,_,edges in G}
    R = {B[v]: r for v,r,_ in G}

    all_open = sum(B.values())

    aa = B['AA']
    states = [(aa, 0, 0)]
    minutes = 26
    for _ in range(minutes):
        print(minutes, len(states))
        minutes -= 1
        next = []
        seen = set()
        states.sort(key=lambda x: -x[2])
        for v, open, reward in states:
            if (v,open) in seen:
                continue
            seen.add((v,open))
            #next.append((v, open, reward))
            if open == all_open:
                next.append((v, open, reward))
            else:
                #next.append((v, open, reward))
                r = R[v]
                if r and not v&open:
                    next.append((v, open|v, reward+r*minutes))
                for e in E[v]:
                    next.append((e, open, reward))
        states = next

    print(max(r for _,_,r in states))

    best = {}
    for _,o,r in states:
        if o in best and best[o] >= r:
            continue
        best[o] = r

    maxpair = []
    def pairs():
        O = sorted(best.keys())
        for i in range(len(O)):
            for j in range(i,len(O)):
                if not O[i] & O[j]:
                    yield(best[O[i]]+best[O[j]])
    print(max(pairs()))

part2_c()