adventofcode/2024/go/utils/dijkstra/graph.go

package dijkstra

import (
	"fmt"
	"math"
	"sync"
)

type ItemGraph struct {
	Nodes []*Node
	Edges map[Node][]*Edge
	lock  sync.RWMutex
}

// AddNode adds a node to the graph
func (g *ItemGraph) AddNode(n *Node) {
	g.lock.Lock()
	g.Nodes = append(g.Nodes, n)
	g.lock.Unlock()
}

// AddEdge adds an edge to the graph
func (g *ItemGraph) AddEdge(n1, n2 *Node, weight int) {
	g.lock.Lock()
	if g.Edges == nil {
		g.Edges = make(map[Node][]*Edge)
	}
	ed1 := Edge{
		Node:   n2,
		Weight: weight,
	}

	ed2 := Edge{
		Node:   n1,
		Weight: weight,
	}
	g.Edges[*n1] = append(g.Edges[*n1], &ed1)
	g.Edges[*n2] = append(g.Edges[*n2], &ed2)
	g.lock.Unlock()
}

// dijkstra implement
func getShortestPath(startNode *Node, endNode *Node, g *ItemGraph) ([]Point, int) {
	visited := make(map[Point]bool)
	dist := make(map[Point]int)
	prev := make(map[Point]Point)
	//pq := make(PriorityQueue, 1)
	//heap.Init(&pq)
	q := NodeQueue{}
	pq := q.NewQ()
	start := Vertex{
		Node:     startNode,
		Distance: 0,
	}
	for _, nval := range g.Nodes {
		dist[nval.Value] = math.MaxInt64
	}
	dist[startNode.Value] = start.Distance
	pq.Enqueue(start)
	//im := 0
	for !pq.IsEmpty() {
		v := pq.Dequeue()
		if visited[v.Node.Value] {
			continue
		}
		visited[v.Node.Value] = true
		near := g.Edges[*v.Node]

		for _, val := range near {
			if !visited[val.Node.Value] {
				if dist[v.Node.Value]+val.Weight < dist[val.Node.Value] {
					store := Vertex{
						Node:     val.Node,
						Distance: dist[v.Node.Value] + val.Weight,
					}
					dist[val.Node.Value] = dist[v.Node.Value] + val.Weight
					//prev[val.Node.Value] = fmt.Sprintf("->%s", v.Node.Value)
					prev[val.Node.Value] = v.Node.Value
					pq.Enqueue(store)
				}
				//visited[val.Node.value] = true
			}
		}
	}
	//	fmt.Println(dist)
	//	fmt.Println(prev)
	pathval := prev[endNode.Value]
	var finalArr []Point
	finalArr = append(finalArr, endNode.Value)
	for pathval != startNode.Value {
		finalArr = append(finalArr, pathval)
		pathval = prev[pathval]
	}
	finalArr = append(finalArr, pathval)
	fmt.Println(finalArr)
	for i, j := 0, len(finalArr)-1; i < j; i, j = i+1, j-1 {
		finalArr[i], finalArr[j] = finalArr[j], finalArr[i]
	}
	return finalArr, dist[endNode.Value]

}

func getAllShortestPaths(startNode *Node, endNode *Node, g *ItemGraph) ([][]Point, int) {
	visited := make(map[Point]bool)
	dist := make(map[Point]int)
	prev := make(map[Point][]Point) // Store multiple predecessors for each node
	q := NodeQueue{}
	pq := q.NewQ()

	start := Vertex{
		Node:     startNode,
		Distance: 0,
	}

	for _, n := range g.Nodes {
		dist[n.Value] = math.MaxInt64
	}
	dist[startNode.Value] = start.Distance
	pq.Enqueue(start)

	// Perform Dijkstra's algorithm
	for !pq.IsEmpty() {
		v := pq.Dequeue()
		if visited[v.Node.Value] {
			continue
		}
		visited[v.Node.Value] = true
		near := g.Edges[*v.Node]

		for _, edge := range near {
			alt := dist[v.Node.Value] + edge.Weight

			// Case 1: Found a shorter path
			if alt < dist[edge.Node.Value] {
				dist[edge.Node.Value] = alt
				prev[edge.Node.Value] = []Point{v.Node.Value} // Reset predecessors
				pq.Enqueue(Vertex{
					Node:     edge.Node,
					Distance: alt,
				})
			}

			// Case 2: Found an equally short path
			if alt == dist[edge.Node.Value] {
				// Add the current node as a valid predecessor if not already present
				found := false
				for _, p := range prev[edge.Node.Value] {
					if p == v.Node.Value {
						found = true
						break
					}
				}
				if !found {
					prev[edge.Node.Value] = append(prev[edge.Node.Value], v.Node.Value)
				}
			}
		}
	}

	// Iteratively reconstruct all paths
	var paths [][]Point
	stack := []struct {
		node Point
		path []Point
	}{
		{node: endNode.Value, path: []Point{}},
	}

	for len(stack) > 0 {
		// Pop the top item from the stack
		current := stack[len(stack)-1]
		stack = stack[:len(stack)-1]

		// Prepend the current node to the path
		newPath := append([]Point{current.node}, current.path...)

		// If we've reached the start node, save the path
		if current.node.X == startNode.Value.X && current.node.Y == startNode.Value.Y {
			paths = append(paths, newPath)
			continue
		}

		// Push all predecessors onto the stack
		for _, predecessor := range prev[current.node] {
			if !Contains(newPath, predecessor) {
				stack = append(stack, struct {
					node Point
					path []Point
				}{
					node: predecessor,
					path: newPath,
				})
			}
		}
	}

	return paths, dist[endNode.Value]
}

func Contains(slice []Point, value Point) bool {
	for _, v := range slice {
		if v == value { // Compare values
			return true
		}
	}
	return false
}

func CreateGraph(data InputGraph) *ItemGraph {
	var g ItemGraph
	nodes := make(map[Point]*Node)
	for _, v := range data.Graph {
		if _, found := nodes[v.Source]; !found {
			nA := Node{v.Source}
			nodes[v.Source] = &nA
			g.AddNode(&nA)
		}
		if _, found := nodes[v.Destination]; !found {
			nA := Node{v.Destination}
			nodes[v.Destination] = &nA
			g.AddNode(&nA)
		}
		g.AddEdge(nodes[v.Source], nodes[v.Destination], v.Weight)
	}
	return &g
}

func GetShortestPath(from, to Point, g *ItemGraph) ([]Point, int) {
	nA := &Node{from}
	nB := &Node{to}

	path, distance := getShortestPath(nA, nB, g)
	return path, distance
}

func GetAllShortestPaths(from, to Point, g *ItemGraph) ([][]Point, int) {
	nA := &Node{from}
	nB := &Node{to}

	paths, distance := getAllShortestPaths(nA, nB, g)
	return paths, distance
}