day 10 part 2 but it's too slow

day10/sol2.go

@@ -0,0 +1,156 @@
+package main
+
+import (
+	"bufio"
+	"container/heap"
+	"fmt"
+	"log"
+	"os"
+	"strconv"
+	"strings"
+)
+
+func die(err error) {
+	log.Fatal(err)
+}
+
+func check(err error) {
+	if err != nil {
+		die(err)
+	}
+}
+
+func main() {
+	solve("sample")
+	solve("input")
+}
+
+func solve(filename string) {
+	input, err := os.Open(filename)
+	check(err)
+
+	scanner := bufio.NewScanner(input)
+	scanner.Split(bufio.ScanLines)
+	total := 0
+	for scanner.Scan() {
+		line := scanner.Text()
+		parts := strings.Fields(line)
+		parts = parts[1:]
+		var jolts []Jolts
+		for _, p := range parts {
+			j, err := parseJolt(p)
+			check(err)
+			jolts = append(jolts, j)
+		}
+		fmt.Printf("%v %v\n", jolts, parts)
+		n := best(jolts[len(jolts)-1], jolts[:len(jolts)-1])
+		fmt.Printf("%s = %v\n", line, n)
+		total += n
+		//fmt.Println(n)
+	}
+	check(scanner.Err())
+	fmt.Println(total)
+
+}
+
+type Jolts [10]int16
+
+func parseJolt(s string) (Jolts, error) {
+	if s[0] == '(' {
+		return parseJoltByIndex(s)
+	} else if s[0] == '{' {
+		return parseJoltByValue(s)
+	}
+	return Jolts{}, fmt.Errorf("invalid jolts: %q", s)
+}
+
+func parseJoltByIndex(s string) (Jolts, error) {
+	var j Jolts
+	s = strings.Trim(s, "()")
+	for _, part := range strings.Split(s, ",") {
+		idx, err := strconv.ParseInt(part, 10, 16)
+		if err != nil {
+			return j, err
+		}
+		j[idx] = 1
+	}
+	return j, nil
+}
+
+func parseJoltByValue(s string) (Jolts, error) {
+	var j Jolts
+	s = strings.Trim(s, "{}")
+	for i, part := range strings.Split(s, ",") {
+		val, err := strconv.ParseInt(part, 10, 16)
+		if err != nil {
+			return j, err
+		}
+		j[i] = int16(val)
+	}
+	return j, nil
+}
+
+// An CostHeap is a min-heap of ints.
+type CostHeap struct {
+	heap []Jolts
+	cost map[Jolts]int
+}
+
+func (h *CostHeap) Len() int           { return len(h.heap) }
+func (h *CostHeap) Less(i, j int) bool { return h.cost[h.heap[i]] < h.cost[h.heap[j]] }
+func (h *CostHeap) Swap(i, j int)      { h.heap[i], h.heap[j] = h.heap[j], h.heap[i] }
+
+func (h *CostHeap) Push(x any) {
+	// Push and Pop use pointer receivers because they modify the slice's length,
+	// not just its contents.
+	h.heap = append(h.heap, x.(Jolts))
+}
+
+func (h *CostHeap) Pop() any {
+	old := h.heap
+	n := len(old)
+	x := old[n-1]
+	h.heap = old[0 : n-1]
+	return x
+}
+
+func best(target Jolts, pool []Jolts) int {
+	var cost = make(map[Jolts]int)
+	h := &CostHeap{nil, cost}
+	heap.Push(h, Jolts{})
+	for h.Len() > 0 {
+		s := heap.Pop(h).(Jolts)
+		if s == target {
+			break
+		}
+		// enumerate all the states that can be reached by toggling button p
+	next:
+		for _, p := range pool {
+			t := s
+			for i, v := range p {
+				t[i] += v
+				if t[i] > target[i] {
+					// blown target, cut path
+					continue next
+				}
+			}
+
+			if cost_t, ok := cost[t]; ok {
+				if cost[s]+1 < cost_t {
+					cost[t] = cost[s] + 1
+					for i, x := range h.heap {
+						if x == t {
+							heap.Fix(h, i)
+							break
+						}
+					}
+				}
+			} else {
+				cost[t] = cost[s] + 1
+				heap.Push(h, t)
+			}
+		}
+	}
+
+	return cost[target]
+}