package main

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

type Edge struct {
	To     int
	Weight int64
}

type item struct {
	vertex   int
	distance int64
}

type priorityQueue []item

func (pq priorityQueue) Len() int            { return len(pq) }
func (pq priorityQueue) Less(i, j int) bool  { return pq[i].distance < pq[j].distance }
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 distances and predecessors from source.
// Unreachable vertices have distance math.MaxInt64 and predecessor -1.
func Dijkstra(graph [][]Edge, source int) ([]int64, []int, error) {
	if source < 0 || source >= len(graph) {
		return nil, nil, fmt.Errorf("source vertex %d is out of range", source)
	}

	dist := make([]int64, len(graph))
	prev := make([]int, len(graph))
	for i := range graph {
		dist[i] = math.MaxInt64
		prev[i] = -1
	}
	dist[source] = 0

	pq := &priorityQueue{{vertex: source, distance: 0}}
	heap.Init(pq)

	for pq.Len() > 0 {
		current := heap.Pop(pq).(item)
		if current.distance != dist[current.vertex] {
			continue // Ignore stale queue entries.
		}

		for _, edge := range graph[current.vertex] {
			if edge.To < 0 || edge.To >= len(graph) {
				return nil, nil, fmt.Errorf("edge points to invalid vertex %d", edge.To)
			}
			if edge.Weight < 0 {
				return nil, nil, fmt.Errorf("negative edge weight from %d to %d", current.vertex, edge.To)
			}
			if current.distance > math.MaxInt64-edge.Weight {
				continue // The candidate distance would overflow.
			}

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

	return dist, prev, nil
}

func main() {
	tests := []struct {
		name   string
		graph  [][]Edge
		source int
		want   []int64
	}{
		{
			name: "shortest routes",
			graph: [][]Edge{
				{{To: 1, Weight: 4}, {To: 2, Weight: 1}},
				{{To: 3, Weight: 1}},
				{{To: 1, Weight: 2}, {To: 3, Weight: 5}},
				{},
			},
			source: 0,
			want:   []int64{0, 3, 1, 4},
		},
		{
			name: "unreachable vertex",
			graph: [][]Edge{
				{{To: 1, Weight: 7}},
				{},
				{},
			},
			source: 0,
			want:   []int64{0, 7, math.MaxInt64},
		},
	}

	for _, test := range tests {
		got, _, err := Dijkstra(test.graph, test.source)
		if err != nil {
			panic(fmt.Sprintf("%s: %v", test.name, err))
		}
		for i := range test.want {
			if got[i] != test.want[i] {
				panic(fmt.Sprintf("%s: vertex %d: got %d, want %d", test.name, i, got[i], test.want[i]))
			}
		}
	}
}