package main

import (
	"container/heap"
	"fmt"
	"math"
)

type Edge struct {
	To     int
	Weight int
}

type item struct {
	node int
	dist int
}

type priorityQueue []item

func (pq priorityQueue) Len() int           { return len(pq) }
func (pq priorityQueue) Less(i, j int) bool { return pq[i].dist < pq[j].dist }
func (pq priorityQueue) Swap(i, j int)      { pq[i], pq[j] = pq[j], pq[i] }

func (pq *priorityQueue) Push(x any) {
	*pq = append(*pq, x.(item))
}

func (pq *priorityQueue) Pop() any {
	old := *pq
	x := old[len(old)-1]
	*pq = old[:len(old)-1]
	return x
}

// Dijkstra returns the minimum distance from source to every vertex.
// Unreachable vertices have distance math.MaxInt.
func Dijkstra(graph [][]Edge, source int) []int {
	if source < 0 || source >= len(graph) {
		panic("source vertex out of range")
	}

	dist := make([]int, len(graph))
	for i := range dist {
		dist[i] = math.MaxInt
	}
	dist[source] = 0

	pq := &priorityQueue{{node: source, dist: 0}}
	heap.Init(pq)

	for pq.Len() > 0 {
		current := heap.Pop(pq).(item)

		// Ignore stale queue entries.
		if current.dist != dist[current.node] {
			continue
		}

		for _, edge := range graph[current.node] {
			if edge.To < 0 || edge.To >= len(graph) {
				panic("edge endpoint out of range")
			}
			if edge.Weight < 0 {
				panic("Dijkstra requires non-negative edge weights")
			}
			if current.dist > math.MaxInt-edge.Weight {
				continue
			}

			candidate := current.dist + edge.Weight
			if candidate < dist[edge.To] {
				dist[edge.To] = candidate
				heap.Push(pq, item{node: edge.To, dist: candidate})
			}
		}
	}

	return dist
}

func assertEqual(name string, got, want []int) {
	if len(got) != len(want) {
		panic(fmt.Sprintf("%s: got %v, want %v", name, got, want))
	}
	for i := range got {
		if got[i] != want[i] {
			panic(fmt.Sprintf("%s: got %v, want %v", name, got, want))
		}
	}
}

func main() {
	graph := [][]Edge{
		{{To: 1, Weight: 4}, {To: 2, Weight: 1}},
		{{To: 3, Weight: 1}},
		{{To: 1, Weight: 2}, {To: 3, Weight: 5}},
		nil,
	}
	assertEqual("shorter indirect path", Dijkstra(graph, 0), []int{0, 3, 1, 4})

	disconnected := [][]Edge{
		{{To: 1, Weight: 7}},
		nil,
		nil,
	}
	assertEqual("unreachable vertex", Dijkstra(disconnected, 0), []int{0, 7, math.MaxInt})
}