dijkstra.go

package graph

import "container/heap"

// edge is a weighted outbound edge.
type edge struct {
	to     int
	weight int
}

// WeightedGraph is a directed graph with non-negative integer weights, stored
// as adjacency lists.
type WeightedGraph struct {
	adj [][]edge
}

// NewWeighted returns a new WeightedGraph with n vertices and no edges.
func NewWeighted(n int) *WeightedGraph {
	if n < 0 {
		n = 0
	}
	return &WeightedGraph{adj: make([][]edge, n)}
}

// Vertices returns the number of vertices.
func (g *WeightedGraph) Vertices() int { return len(g.adj) }

// AddEdge adds a directed edge u -> v with the given weight. Panics if the
// weight is negative (Dijkstra requires non-negative weights).
func (g *WeightedGraph) AddEdge(u, v, w int) {
	if w < 0 {
		panic("graph: negative weight not allowed")
	}
	g.adj[u] = append(g.adj[u], edge{to: v, weight: w})
}

// Inf is the value used in dist[] for unreachable vertices.
const Inf = int(1 << 62)

// pqItem is one entry in the priority queue used by Dijkstra.
type pqItem struct {
	node  int
	dist  int
	index int
}

// pq implements heap.Interface, smallest distance first.
type pq []*pqItem

func (p pq) Len() int           { return len(p) }
func (p pq) Less(i, j int) bool { return p[i].dist < p[j].dist }
func (p pq) Swap(i, j int) {
	p[i], p[j] = p[j], p[i]
	p[i].index, p[j].index = i, j
}
func (p *pq) Push(x any) {
	it := x.(*pqItem)
	it.index = len(*p)
	*p = append(*p, it)
}
func (p *pq) Pop() any {
	old := *p
	n := len(old)
	it := old[n-1]
	old[n-1] = nil
	it.index = -1
	*p = old[:n-1]
	return it
}

// Dijkstra computes the shortest distance from src to every other vertex
// along with a `prev` array suitable for path reconstruction. `dist[v]`
// equals Inf if v is unreachable from src; `prev[v]` is -1 for the source
// or unreachable vertices.
//
// Time:  O((V + E) log V) using a binary heap.
// Space: O(V).
func (g *WeightedGraph) Dijkstra(src int) (dist, prev []int) {
	n := len(g.adj)
	dist = make([]int, n)
	prev = make([]int, n)
	for i := range dist {
		dist[i] = Inf
		prev[i] = -1
	}
	if src < 0 || src >= n {
		return
	}
	dist[src] = 0

	q := &pq{}
	heap.Init(q)
	heap.Push(q, &pqItem{node: src, dist: 0})

	for q.Len() > 0 {
		it := heap.Pop(q).(*pqItem)
		u := it.node
		if it.dist != dist[u] {
			// stale entry left over from a relaxation
			continue
		}
		for _, e := range g.adj[u] {
			nd := dist[u] + e.weight
			if nd < dist[e.to] {
				dist[e.to] = nd
				prev[e.to] = u
				heap.Push(q, &pqItem{node: e.to, dist: nd})
			}
		}
	}
	return dist, prev
}

// ReconstructPath returns the path src -> dst given the prev array from
// Dijkstra. Returns nil if dst is unreachable.
func ReconstructPath(prev []int, src, dst int) []int {
	if dst < 0 || dst >= len(prev) {
		return nil
	}
	if dst != src && prev[dst] == -1 {
		return nil
	}
	path := []int{}
	for at := dst; at != -1; at = prev[at] {
		path = append(path, at)
		if at == src {
			break
		}
	}
	if path[len(path)-1] != src {
		return nil
	}
	// reverse in place
	for i, j := 0, len(path)-1; i < j; i, j = i+1, j-1 {
		path[i], path[j] = path[j], path[i]
	}
	return path
}