Merge pull request 'chore: fix dijkstra' (#17) from day17-enhance into main

Reviewed-on: #17

aoc_test.go

@@ -1,6 +1,8 @@
 package aoc_test
 
 import (
+	"fmt"
+	"sort"
 	"testing"
 
 	"github.com/matryer/is"
@@ -79,15 +81,130 @@ func TestList(t *testing.T) {
 	is.Equal(a, 5)
 }
 
-func TestGraph(t *testing.T) {
-	g := aoc.Graph[int, uint](7)
-	g.AddEdge(0, 1, 2)
-	g.AddEdge(0, 2, 6)
-	g.AddEdge(1, 3, 5)
-	g.AddEdge(2, 3, 8)
-	g.AddEdge(3, 4, 10)
-	g.AddEdge(3, 5, 15)
-	g.AddEdge(4, 6, 2)
-	g.AddEdge(5, 6, 6)
-	// g.Dijkstra(0)
+func TestPriorityQueue(t *testing.T) {
+	is := is.New(t)
+
+	type elem [2]int
+	less := func(a, b elem) bool {
+		return b[0] < a[0]
+	}
+
+	pq := aoc.PriorityQueue(less)
+
+	pq.Enqueue(elem{1, 4})
+	pq.Enqueue(elem{3, 2})
+	pq.Enqueue(elem{2, 3})
+	pq.Enqueue(elem{4, 1})
+
+	v, ok := pq.Dequeue()
+	is.True(ok)
+	is.Equal(v, elem{4, 1})
+
+	v, ok = pq.Dequeue()
+	is.True(ok)
+	is.Equal(v, elem{3, 2})
+
+	v, ok = pq.Dequeue()
+	is.True(ok)
+	is.Equal(v, elem{2, 3})
+
+	v, ok = pq.Dequeue()
+	is.True(ok)
+	is.Equal(v, elem{1, 4})
+
+	v, ok = pq.Dequeue()
+	is.True(!ok)
+	is.Equal(v, elem{})
+}
+
+func TestSet(t *testing.T) {
+	is := is.New(t)
+
+	s := aoc.Set(1, 2, 3)
+	is.True(!s.Has(0))
+	is.True(s.Has(1))
+	is.True(s.Has(2))
+	is.True(s.Has(3))
+	is.True(!s.Has(4))
+
+	s.Add(4)
+	is.True(s.Has(4))
+
+	items := s.Items()
+	sort.Ints(items)
+	is.Equal(items, []int{1, 2, 3, 4})
+}
+
+// func TestGraph(t *testing.T) {
+// 	g := aoc.Graph[int, uint](7)
+// 	g.AddEdge(0, 1, 2)
+// 	g.AddEdge(0, 2, 6)
+// 	g.AddEdge(1, 3, 5)
+// 	g.AddEdge(2, 3, 8)
+// 	g.AddEdge(3, 4, 10)
+// 	g.AddEdge(3, 5, 15)
+// 	g.AddEdge(4, 6, 2)
+// 	g.AddEdge(5, 6, 6)
+// 	// g.Dijkstra(0)
+// }
+
+func ExamplePriorityQueue() {
+	type memo struct {
+		pt    int
+		score int
+	}
+	less := func(a, b memo) bool { return a.score < b.score }
+
+	adj := map[int][][2]int{
+		0: {{1, 2}, {2, 6}},
+		1: {{3, 5}},
+		2: {{3, 8}},
+		3: {{4, 10}, {5, 15}},
+		4: {{6, 2}},
+		5: {{6, 6}},
+	}
+
+	pq := aoc.PriorityQueue(less)
+	visited := aoc.Set([]int{}...)
+	dist := aoc.DefaultMap[int](int(^uint(0) >> 1))
+
+	dist.Set(0, 0)
+	pq.Enqueue(memo{0, 0})
+
+	for !pq.IsEmpty() {
+		m, _ := pq.Dequeue()
+
+		u := m.pt
+		if visited.Has(u) {
+			continue
+		}
+		visited.Add(u)
+
+		du, _ := dist.Get(u)
+
+		for _, edge := range adj[u] {
+			v, w := edge[0], edge[1]
+			dv, _ := dist.Get(v)
+
+			if !visited.Has(v) && du+w < dv {
+				dist.Set(v, du+w)
+				pq.Enqueue(memo{v, du + w})
+			}
+		}
+	}
+
+	items := dist.Items()
+	sort.Slice(items, func(i, j int) bool { return items[i].K < items[j].K })
+	for _, v := range items {
+		fmt.Printf("point %d is %d steps away.\n", v.K, v.V)
+	}
+
+	// Output:
+	// point 0 is 0 steps away.
+	// point 1 is 2 steps away.
+	// point 2 is 6 steps away.
+	// point 3 is 7 steps away.
+	// point 4 is 17 steps away.
+	// point 5 is 22 steps away.
+	// point 6 is 19 steps away.
 }

day01/main_test.go

@@ -30,7 +30,6 @@ func TestExample1(t *testing.T) {
 	is.Equal(result.sum, 142)
 }
 
-
 func TestExample2(t *testing.T) {
 	is := is.New(t)
 	scan := bufio.NewScanner(bytes.NewReader(example2))

day14/main_test.go

@@ -43,4 +43,3 @@ func TestSolution(t *testing.T) {
 	is.Equal(result.valuePT1, 110407)
 	is.Equal(result.valuePT2, 87273)
 }
-

day16/main.go

@@ -96,7 +96,6 @@ func (m *Map) Get(p [2]int) rune {
 	return (*m)[p[0]][p[1]]
 }
 
-
 func runCycle(m Map, r ray) int {
 	current := r
 

day17/main.go

@@ -4,9 +4,9 @@ import (
 	"bufio"
 	_ "embed"
 	"fmt"
-	"sort"
 
 	aoc "go.sour.is/advent-of-code"
+	"golang.org/x/exp/maps"
 )
 
 // var log = aoc.Log
@@ -21,7 +21,7 @@ type result struct {
 func (r result) String() string { return fmt.Sprintf("%#v", r) }
 
 func run(scan *bufio.Scanner) (*result, error) {
-	var m Map
+	var m aoc.Map[rune]
 
 	for scan.Scan() {
 		text := scan.Text()
@@ -35,180 +35,97 @@ func run(scan *bufio.Scanner) (*result, error) {
 	return &result, nil
 }
 
-var (
-	ZERO = point{0, 0}
+func search(m aoc.Map[rune], minSteps, maxSteps int) int {
+	type direction int8
+	type rotate int8
 
-	UP = point{-1, 0}
-	DN = point{1, 0}
-	LF = point{0, -1}
-	RT = point{0, 1}
-
-	INF = int(^uint(0) >> 1)
+	const (
+		CW  rotate = 1
+		CCW rotate = -1
 	)
 
-type Map [][]rune
+	var (
+		U = aoc.Point{-1, 0}
+		R = aoc.Point{0, 1}
+		D = aoc.Point{1, 0}
+		L = aoc.Point{0, -1}
+	)
 
-func (m *Map) Get(p point) (point, int, bool) {
-	if !m.Valid(p) {
-		return [2]int{0, 0}, 0, false
+	var Directions = map[aoc.Point]direction{
+		U: 0, // U
+		R: 1, // R
+		D: 2, // D
+		L: 3, // L
 	}
+	var DirectionIDX = maps.Keys(Directions)
 
-	return p, int((*m)[p[0]][p[1]] - '0'), true
-}
-func (m *Map) GetNeighbor(p point, d point) (point, int, bool) {
-	return m.Get(p.add(d))
-}
-func (m *Map) Size() (int, int) {
-	if m == nil || len(*m) == 0 {
-		return 0, 0
-	}
-	return len(*m), len((*m)[0])
-}
-func (m *Map) Neighbors(p point) []point {
-	var lis []point
-	for _, d := range []point{UP, DN, LF, RT} {
-		if p, _, ok := m.GetNeighbor(p, d); ok {
-			lis = append(lis, p)
-		}
-	}
-	return lis
-}
-func (m *Map) NeighborDirections(p point) []point {
-	var lis []point
-	for _, d := range []point{UP, DN, LF, RT} {
-		if m.Valid(p.add(d)) {
-			lis = append(lis, d)
-		}
-	}
-	return lis
-}
-func (m *Map) Valid(p point) bool {
 	rows, cols := m.Size()
-	return p[0] >= 0 && p[0] < rows && p[1] >= 0 && p[1] < cols
+	target := aoc.Point{rows - 1, cols - 1}
+
+	type position struct {
+		loc       aoc.Point
+		direction aoc.Point
+		steps     int
+	}
+
+	step := func(p position) position {
+		return position{p.loc.Add(p.direction), p.direction, p.steps + 1}
+	}
+	rotateAndStep := func(p position, towards rotate) position {
+		d := DirectionIDX[(int8(Directions[p.direction])+int8(towards)+4)%4]
+		return position{p.loc.Add(d), d, 1}
 	}
 
 	type memo struct {
-	h int
-	s int
-	p point
-	d point
+		cost int
+		position
+	}
+	less := func(a, b memo) bool {
+		if a.cost != b.cost {
+			return a.cost < b.cost
+		}
+		if a.position.loc != b.position.loc {
+			return b.position.loc.Less(a.position.loc)
+		}
+		if a.position.direction != b.position.direction {
+			return b.position.direction.Less(a.position.direction)
+		}
+		return b.steps < a.steps
 	}
 
-func (memo) sort(a, b memo) bool {
-	if a.h != b.h {
-		return a.h < b.h
-	}
-
-	if a.s != b.s {
-		return a.s < b.s
-	}
-
-	if a.p != b.p {
-		return a.p.less(b.p)
-	}
-
-	return a.d.less(b.d)
-}
-
-type priorityQueue[T any, U []T] struct {
-	elems U
-	sort  func(a, b T) bool
-}
-
-func PriorityQueue[T any, U []T](sort func(a, b T) bool) *priorityQueue[T, U] {
-	return &priorityQueue[T, U]{sort: sort}
-}
-func (pq *priorityQueue[T, U]) Enqueue(elem T) {
-	pq.elems = append(pq.elems, elem)
-	sort.Slice(pq.elems, func(i, j int) bool { return pq.sort(pq.elems[i], pq.elems[j]) })
-}
-func (pq *priorityQueue[T, I]) IsEmpty() bool {
-	return len(pq.elems) == 0
-}
-func (pq *priorityQueue[T, I]) Dequeue() (T, bool) {
-	var elem T
-	if pq.IsEmpty() {
-		return elem, false
-	}
-
-	elem, pq.elems = pq.elems[0], pq.elems[1:]
-	return elem, true
-}
-
-func heuristic(m Map, p point) int {
-	rows, cols := m.Size()
-	return rows - p[0] + cols - p[1]
-}
-
-func search(m Map, minSize, maxSize int) int {
-	rows, cols := m.Size()
-	END := point{rows - 1, cols - 1}
-
-	visited := make(map[vector]int)
-	pq := PriorityQueue(memo{}.sort)
-	pq.Enqueue(memo{h: heuristic(m, point{0, 0}), p: point{0, 0}, d: DN})
+	pq := aoc.PriorityQueue(less)
+	pq.Enqueue(memo{position: position{direction: D}})
+	pq.Enqueue(memo{position: position{direction: R}})
+	visited := aoc.Set[position]()
 
 	for !pq.IsEmpty() {
-		mem, _ := pq.Dequeue()
-		fmt.Println(mem)
-		if mem.h > dmap(visited, vector{mem.p[0], mem.p[1], mem.d[0], mem.d[1]}, INF) {
+		current, _ := pq.Dequeue()
+
+		if current.loc == target && current.steps >= minSteps {
+			return current.cost
+		}
+
+		seen := position{loc: current.loc, steps: current.steps}
+		if visited.Has(seen) {
 			continue
 		}
+		visited.Add(seen)
 
-		if mem.p == END {
-			return mem.s
+		if left := rotateAndStep(current.position, CCW); current.steps >= minSteps && m.Valid(left.loc) {
+			_, cost, _ := m.Get(left.loc)
+			pq.Enqueue(memo{cost: current.cost + int(cost-'0'), position: left})
 		}
 
-		for _, nd := range m.NeighborDirections(mem.p) {
-			if nd[0] == 0 && mem.d == RT || nd[1] == 0 && mem.d == DN {
-				continue
+		if right := rotateAndStep(current.position, CW); current.steps >= minSteps && m.Valid(right.loc) {
+			_, cost, _ := m.Get(right.loc)
+			pq.Enqueue(memo{cost: current.cost + int(cost-'0'), position: right})
 		}
 
-			dscore := 0
-
-			for _, size := range irange(1, maxSize+1) {
-				np := mem.p.add(nd.scale(size))
-				_, s, ok := m.Get(np)
-
-				if !ok {
-					break
-				}
-
-				dscore += s
-				pscore := mem.s + dscore
-
-				nh := heuristic(m, np) + pscore
-				vec := vector{np[0], np[1], nd[0], nd[1]}
-
-				if size >= minSize && nh < dmap(visited, vec, INF) {
-					pq.Enqueue(memo{nh, pscore, np, nd})
-					visited[vec] = nh
-				}
-			}
+		if forward := step(current.position); current.steps < maxSteps && m.Valid(forward.loc) {
+			_, cost, _ := m.Get(forward.loc)
+			pq.Enqueue(memo{cost: current.cost + int(cost-'0'), position: forward})
 		}
 	}
 
-	return 0
+	return -1
 }
-
-func dmap[K comparable, V any](m map[K]V, k K, d V) V {
-	if v, ok := m[k]; ok {
-		return v
-	}
-	return d
-}
-func irange(a, b int) []int {
-	lis := make([]int, b-a)
-	for i := range lis {
-		lis[i] = i + a
-	}
-	return lis
-}
-
-type point [2]int
-
-func (p point) add(a point) point { return point{p[0] + a[0], p[1] + a[1]} }
-func (p point) scale(m int) point { return point{p[0] * m, p[1] * m} }
-func (p point) less(a point) bool { return p[0] < a[0] || p[1] < a[1] }
-
-type vector [4]int

day18/main.go

@@ -87,4 +87,3 @@ func findArea(vecs []aoc.Vector) int {
 
 	return aoc.NumPoints(shoelace, borderLength)
 }
-

grids.go

@@ -17,6 +17,12 @@ func (p Point) Add(a Point) Point {
 func (p Point) Scale(m int) Point {
 	return Point{p[0] * m, p[1] * m}
 }
+func (p Point) Less(b Point) bool {
+	if p[0] != b[0] {
+		return p[0] < b[0]
+	}
+	return p[1] < b[1]
+}
 
 func Transpose[T any](matrix [][]T) [][]T {
 	rows, cols := len(matrix), len(matrix[0])
@@ -49,3 +55,24 @@ func NumPoints(outline []Point, borderLength int) int {
 	// pick's theorem - find the number of points in a shape given its area
 	return (ABS(area) - borderLength/2 + 1) + borderLength
 }
+
+type Map[T any] [][]T
+
+func (m *Map[T]) Get(p Point) (Point, T, bool) {
+	var zero T
+	if !m.Valid(p) {
+		return [2]int{0, 0}, zero, false
+	}
+
+	return p, (*m)[p[0]][p[1]], true
+}
+func (m *Map[T]) Size() (int, int) {
+	if m == nil || len(*m) == 0 {
+		return 0, 0
+	}
+	return len(*m), len((*m)[0])
+}
+func (m *Map[T]) Valid(p Point) bool {
+	rows, cols := m.Size()
+	return p[0] >= 0 && p[0] < rows && p[1] >= 0 && p[1] < cols
+}

itertools.go

@@ -31,7 +31,6 @@ func Reverse[T any](arr []T) []T {
 	return arr
 }
 
-
 func SliceMap[T, U any](fn func(T) U, in ...T) []U {
 	lis := make([]U, len(in))
 	for i := range lis {

lists.go

@@ -2,7 +2,6 @@ package aoc
 
 import "fmt"
 
-
 type Node[T any] struct {
 	value T
 	pos   int

runner.go

@@ -41,7 +41,6 @@ func Logf(format string, v ...any) {
 	fmt.Fprintf(os.Stderr, format, v...)
 }
 
-
 func ReadStringToInts(fields []string) []int {
 	return SliceMap(Atoi, fields...)
 }

search.go

@@ -4,100 +4,32 @@ import (
 	"sort"
 )
 
-type PQElem[T any, I integer] struct {
-	Value    T
-	Priority I
-}
-type PQList[T any, I integer] []PQElem[T, I]
-
-func (pq PQList[T, I]) Len() int {
-	return len(pq)
-}
-func (pq PQList[T, I]) Less(i int, j int) bool {
-	return pq[i].Priority < pq[j].Priority
-}
-func (pq PQList[T, I]) Swap(i int, j int) {
-	pq[i], pq[j] = pq[j], pq[i]
+type priorityQueue[T any, U []T] struct {
+	elems U
+	less  func(a, b T) bool
 }
 
-var _ sort.Interface = (*PQList[rune, int])(nil)
-
-type PriorityQueue[T any, I integer] struct {
-	elem PQList[T, I]
+func PriorityQueue[T any, U []T](less func(a, b T) bool) *priorityQueue[T, U] {
+	return &priorityQueue[T, U]{less: less}
 }
-
-func (pq *PriorityQueue[T, I]) Enqueue(elem T, priority I) {
-	pq.elem = append(pq.elem, PQElem[T, I]{elem, priority})
-	sort.Sort(pq.elem)
+func (pq *priorityQueue[T, U]) Enqueue(elem T) {
+	pq.elems = append(pq.elems, elem)
+	sort.Slice(pq.elems, func(i, j int) bool { return pq.less(pq.elems[j], pq.elems[i]) })
 }
-func (pq *PriorityQueue[T, I]) IsEmpty() bool {
-	return len(pq.elem) == 0
+func (pq *priorityQueue[T, I]) IsEmpty() bool {
+	return len(pq.elems) == 0
 }
-func (pq *PriorityQueue[T, I]) Dequeue() (T, bool) {
+func (pq *priorityQueue[T, I]) Dequeue() (T, bool) {
 	var elem T
 	if pq.IsEmpty() {
 		return elem, false
 	}
 
-	elem, pq.elem = pq.elem[0].Value, pq.elem[1:]
+	pq.elems, elem = pq.elems[:len(pq.elems)-1], pq.elems[len(pq.elems)-1]
 	return elem, true
 }
 
-type Vertex[V comparable, I integer] struct {
-	to    V
-	score I
-}
-type graph[V comparable, I uinteger] struct {
-	adj map[V][]Vertex[V, I]
-}
-
-func Graph[V comparable, I uinteger](size int) *graph[V, I] {
-	return &graph[V, I]{
-		adj: make(map[V][]Vertex[V, I], size),
-	}
-}
-func (g *graph[V, I]) AddEdge(u, v V, w I) {
-	g.adj[u] = append(g.adj[u], Vertex[V, I]{to: v, score: w})
-	g.adj[v] = append(g.adj[v], Vertex[V, I]{to: u, score: w})
-}
-func (g *graph[V, I]) Dijkstra(m interface{ Get() }, src V) map[V]I {
-	pq := PriorityQueue[V, I]{}
-	dist := make(map[V]I, len(g.adj))
-	visited := make(map[V]bool, len(g.adj))
-	var INF I
-	INF = ^INF
-
-	pq.Enqueue(src, 0)
-	dist[src] = 0
-
-	for !pq.IsEmpty() {
-		u, _ := pq.Dequeue()
-
-		if _, ok := visited[u]; ok {
-			continue
-		}
-		visited[u] = true
-
-		for _, v := range g.adj[u] {
-			_, ok := visited[v.to]
-			var du, dv I
-			if d, inf := dist[u]; !inf {
-				du = INF
-			} else {
-				du = d
-			}
-			if d, inf := dist[v.to]; !inf {
-				dv = INF
-			} else {
-				dv = d
-			}
-
-			if !ok && du+v.score < dv {
-				dist[v.to] = du + v.score
-				pq.Enqueue(v.to, du+v.score)
-			}
-		}
-	}
-
-	return dist
+type DS[T comparable] struct {
+	*priorityQueue[T, []T]
+	*set[T]
 }

set.go

@@ -0,0 +1,59 @@
+package aoc
+
+import "golang.org/x/exp/maps"
+
+type set[T comparable] map[T]struct{}
+
+func Set[T comparable](arr ...T) set[T] {
+	m := make(set[T], len(arr))
+	for _, a := range arr {
+		m[a] = struct{}{}
+	}
+	return m
+}
+func (m *set[T]) Add(a T) {
+	(*m)[a] = struct{}{}
+}
+func (m *set[T]) Items() []T {
+	return maps.Keys(*m)
+}
+func (m *set[T]) Has(a T) bool {
+	var ok bool
+	_, ok = (*m)[a]
+	return ok
+}
+
+type defaultMap[K comparable, V any] struct {
+	m map[K]V
+	d V
+}
+
+func DefaultMap[K comparable, V any](d V) *defaultMap[K, V] {
+	return &defaultMap[K, V]{
+		make(map[K]V),
+		d,
+	}
+}
+
+func (m *defaultMap[K, V]) Set(k K, v V) {
+	m.m[k] = v
+}
+func (m *defaultMap[K, V]) Get(k K) (V, bool) {
+	if v, ok := m.m[k]; ok {
+		return v, true
+	}
+	return m.d, false
+}
+
+type pair[K, V any] struct {
+	K K
+	V V
+}
+
+func (m *defaultMap[K, V]) Items() []pair[K, V] {
+	var items = make([]pair[K, V], 0, len(m.m))
+	for k, v := range m.m {
+		items = append(items, pair[K, V]{k, v})
+	}
+	return items
+}