adventofcode/2022/day16/main.go

package main

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"

	h "git.bullercodeworks.com/brian/adventofcode/helpers"
)

var valveregexp = regexp.MustCompile(`([A-Z]{2}).*=(\d+);.*?((?:[A-Z]{2}(?:, )?)+)`)

func main() {
	rooms := parseRooms(h.StdinToStringSlice())
	graph := floydWarshall(rooms)

	v := NewVolcano(rooms, graph)
	r := part1(v)
	v = NewVolcano(rooms, graph)
	part2(v, r)
}

func part1(v *Volcano) uint16 {
	var dfs func(target, pressure, minute, on, node uint16) uint16
	dfs = func(target, pressure, minute, on, node uint16) uint16 {
		max := pressure
		for _, w := range v.goodbits {
			if node == w[0] || w[0] == v.start || w[0]&on != 0 {
				continue
			}
			l := v.bitgraphs[node|w[0]] + 1
			if minute+l > target {
				continue
			}
			if next := dfs(target, pressure+(target-minute-l)*w[1], minute+l, on|w[0], w[0]); next > max {
				max = next
			}
		}
		return max
	}

	res := dfs(30, 0, 0, 0, v.start)
	fmt.Println("# Part 1")
	fmt.Println(res)
	return res
}

func part2(v *Volcano, p1 uint16) {
	var dfspaths func(target, pressure, minute, on, node, path uint16) [][2]uint16
	dfspaths = func(target, pressure, minute, on, node, path uint16) [][2]uint16 {
		paths := [][2]uint16{{pressure, path}}
		for _, w := range v.goodbits {
			if w[0] == node || w[0] == v.start || w[0]&on != 0 {
				continue
			}
			l := v.bitgraphs[node|w[0]] + 1
			if minute+l > target {
				continue
			}
			paths = append(paths, dfspaths(target, pressure+(target-minute-l)*w[1], minute+l, on|w[0], w[0], path|w[0])...)
		}
		return paths
	}

	allpaths := dfspaths(26, 0, 0, 0, v.start, 0)

	// reduce paths (presumably, both paths are at least half of part 1)
	var trimpaths [][2]uint16
	for _, p := range allpaths {
		if p[0] > p1/2 {
			trimpaths = append(trimpaths, p)
		}
	}

	// compare all paths to find max
	var max uint16 = 0
	for idx := 0; idx < len(trimpaths); idx += 1 {
		for jdx := idx + 1; jdx < len(trimpaths); jdx += 1 {
			if trimpaths[idx][1]&trimpaths[jdx][1] != 0 {
				continue
			}
			if m := trimpaths[idx][0] + trimpaths[jdx][0]; m > max {
				max = m
			}
		}
	}

	fmt.Println("# Part 2")
	fmt.Println(max)
}

type Volcano struct {
	start uint16

	goodrooms []*Room
	goodbits  [][2]uint16

	bitfield  map[*Room]uint16
	bitgraphs []uint16
}

func NewVolcano(rooms []*Room, graph map[*Room]map[*Room]uint16) *Volcano {
	v := Volcano{}

	// pick valves with flow and starting point
	for _, r := range rooms {
		if r.Rate > 0 || r.Name == "AA" {
			v.goodrooms = append(v.goodrooms, r)
		}
	}

	// assign bits
	v.bitfield = make(map[*Room]uint16)
	for idx, r := range v.goodrooms {
		v.bitfield[r] = 1 << idx
	}

	// find start
	for _, r := range v.goodrooms {
		if r.Name == "AA" {
			v.start = v.bitfield[r]
			break
		}
	}

	// create slice for fast edge lookup
	v.bitgraphs = make([]uint16, 0xffff)
	for _, v1 := range v.goodrooms {
		for _, v2 := range v.goodrooms {
			v.bitgraphs[v.bitfield[v1]|v.bitfield[v2]] = graph[v1][v2]
		}
	}

	// create slice for fast node lookup
	v.goodbits = make([][2]uint16, len(v.goodrooms))
	for idx, r := range v.goodrooms {
		v.goodbits[idx] = [2]uint16{v.bitfield[r], r.Rate}
	}
	return &v
}

type Room struct {
	Name  string
	Rate  uint16
	Edges string
}

func parseRooms(inp []string) []*Room {
	rooms := make([]*Room, len(inp))
	for idx, line := range inp {
		m := valveregexp.FindStringSubmatch(line)
		i, _ := strconv.Atoi(m[2])
		rooms[idx] = &Room{Name: m[1], Rate: uint16(i), Edges: m[3]}
	}
	return rooms
}

func floydWarshall(valves []*Room) map[*Room]map[*Room]uint16 {
	graph := make(map[*Room]map[*Room]uint16)
	for _, v1 := range valves {
		graph[v1] = make(map[*Room]uint16)
		for _, v2 := range valves {
			if v1 == v2 {
				graph[v1][v2] = 0
			} else if strings.Contains(v1.Edges, v2.Name) {
				graph[v1][v2] = 1
			} else {
				graph[v1][v2] = 0xff
			}
		}
	}

	for _, k := range valves {
		for _, i := range valves {
			for _, j := range valves {
				if graph[i][j] > graph[i][k]+graph[k][j] {
					graph[i][j] = graph[i][k] + graph[k][j]
				}
			}
		}
	}

	return graph
}