loading...

Implementing Dijkstra Algorithm in Go

jonatasbaldin profile image Jonatas Baldin ・5 min read

Originally posted at Deploy Everyday.


I didn't go through Computer Science, had a very early dropout. Since I started working in the tech space, I cared a lot about RESTful APIs, preventing systems to fail. Algorithms and data structures always seemed daunting and terrifying.

After picking up the Grokking Algorithms book and giving up three times, I finally decided to put some real effort into it. And it's being an amazing adventure, I can't recommend it enough!

Chapter 7 is all about the famous Dijkstra Algorithm: finding the shortest path for a vertice in a graph. A bunch of pomodoros later, I got it. The code below is as commented as possible, to solidify the knowledge in my head and help others as well. I also wrote unit tests and the answers to the book's exercises here.

I won't go in much detail explaining the theory behind it, other resources do a better job I'd be capable of. Start with this quick video, take a look at Brilliant's article and, of course, read the already mentioned book if you can.

It probably can be improved, as everything can. I think this is a nice balance between "working" and "readable". If you have any nice tips to contribute, please leave it in the comments ✨

Well, here it is, the Dijkstra Algorithm in Go:

package main

import (
    "fmt"
    "sort"
    "strconv"
)

type Graph struct {
    Edges []*Edge
    Nodes []*Node
}

type Edge struct {
    Parent *Node
    Child  *Node
    Cost   int
}

type Node struct {
    Name string
}

const Infinity = int(^uint(0) >> 1)

// AddEdge adds an Edge to the Graph
func (g *Graph) AddEdge(parent, child *Node, cost int) {
    edge := &Edge{
        Parent: parent,
        Child:  child,
        Cost:   cost,
    }

    g.Edges = append(g.Edges, edge)
    g.AddNode(parent)
    g.AddNode(child)
}

// AddNode adds a Node to the Graph list of Nodes, if the the node wasn't already added
func (g *Graph) AddNode(node *Node) {
    var isPresent bool
    for _, n := range g.Nodes {
        if n == node {
            isPresent = true
        }
    }

    if !isPresent {
        g.Nodes = append(g.Nodes, node)
    }
}

// String returns a string representation of the Graph
func (g *Graph) String() string {
    var s string

    s += "Edges:\n"
    for _, edge := range g.Edges {
        s += edge.Parent.Name + " -> " + edge.Child.Name + " = " + strconv.Itoa(edge.Cost)
        s += "\n"
    }
    s += "\n"

    s += "Nodes: "
    for i, node := range g.Nodes {
        if i == len(g.Nodes)-1 {
            s += node.Name
        } else {
            s += node.Name + ", "
        }
    }
    s += "\n"

    return s
}

// Dijkstra implements THE Dijkstra algorithm
// Returns the shortest path from startNode to all the other Nodes
func (g *Graph) Dijkstra(startNode *Node) (shortestPathTable string) {

    // First, we instantiate a "Cost Table", it will hold the information:
    // "From startNode, what's is the cost to all the other Nodes?"
    // When initialized, It looks like this:
    // NODE  COST
    //  A     0    // The startNode has always the lowest cost to itself, in this case, 0
    //  B    Inf   // the distance to all the other Nodes are unknown, so we mark as Infinity
    //  C    Inf
    // ...
    costTable := g.NewCostTable(startNode)

    // An empty list of "visited" Nodes. Everytime the algorithm runs on a Node, we add it here
    var visited []*Node

    // A loop to visit all Nodes
    for len(visited) != len(g.Nodes) {

        // Get closest non visited Node (lower cost) from the costTable
        node := getClosestNonVisitedNode(costTable, visited)

        // Mark Node as visited
        visited = append(visited, node)

        // Get Node's Edges (its neighbors)
        nodeEdges := g.GetNodeEdges(node)

        for _, edge := range nodeEdges {

            // The distance to that neighbor, let's say B is the cost from the costTable + the cost to get there (Edge cost)
            // In the first run, the costTable says it's "Infinity"
            // Plus the actual cost, let's say "5"
            // The distance becomes "5"
            distanceToNeighbor := costTable[node] + edge.Cost

            // If the distance above is lesser than the distance currently in the costTable for that neighbor
            if distanceToNeighbor < costTable[edge.Child] {

                // Update the costTable for that neighbor
                costTable[edge.Child] = distanceToNeighbor
            }
        }
    }

    // Make the costTable nice to read :)
    for node, cost := range costTable {
        shortestPathTable += fmt.Sprintf("Distance from %s to %s = %d\n", startNode.Name, node.Name, cost)
    }

    return shortestPathTable
}

// NewCostTable returns an initialized cost table for the Dijkstra algorithm work with
// by default, the lowest cost is assigned to the startNode – so the algorithm starts from there
// all the other Nodes in the Graph receives the Infinity value
func (g *Graph) NewCostTable(startNode *Node) map[*Node]int {
    costTable := make(map[*Node]int)
    costTable[startNode] = 0

    for _, node := range g.Nodes {
        if node != startNode {
            costTable[node] = Infinity
        }
    }

    return costTable
}

// GetNodeEdges returns all the Edges that start with the specified Node
// In other terms, returns all the Edges connecting to the Node's neighbors
func (g *Graph) GetNodeEdges(node *Node) (edges []*Edge) {
    for _, edge := range g.Edges {
        if edge.Parent == node {
            edges = append(edges, edge)
        }
    }

    return edges
}

// getClosestNonVisitedNode returns the closest Node (with the lower cost) from the costTable
// **if the node hasn't been visited yet**
func getClosestNonVisitedNode(costTable map[*Node]int, visited []*Node) *Node {
    type CostTableToSort struct {
        Node *Node
        Cost int
    }
    var sorted []CostTableToSort

    // Verify if the Node has been visited already
    for node, cost := range costTable {
        var isVisited bool
        for _, visitedNode := range visited {
            if node == visitedNode {
                isVisited = true
            }
        }
        // If not, add them to the sorted slice
        if !isVisited {
            sorted = append(sorted, CostTableToSort{node, cost})
        }
    }

    // We need the Node with the lower cost from the costTable
    // So it's important to sort it
    // Here I'm using an anonymous struct to make it easier to sort a map
    sort.Slice(sorted, func(i, j int) bool {
        return sorted[i].Cost < sorted[j].Cost
    })

    return sorted[0].Node
}

func main() {
    a := &Node{Name: "a"}
    b := &Node{Name: "b"}
    c := &Node{Name: "c"}
    d := &Node{Name: "d"}
    e := &Node{Name: "e"}
    f := &Node{Name: "f"}
    g := &Node{Name: "g"}

    graph := Graph{}
    graph.AddEdge(a, c, 2)
    graph.AddEdge(a, b, 5)
    graph.AddEdge(c, b, 1)
    graph.AddEdge(c, d, 9)
    graph.AddEdge(b, d, 4)
    graph.AddEdge(d, e, 2)
    graph.AddEdge(d, g, 30)
    graph.AddEdge(d, f, 10)
    graph.AddEdge(f, g, 1)

    fmt.Println(graph.Dijkstra(a))
}

The output will be:

Distance from a to a = 0
Distance from a to c = 2
Distance from a to b = 3
Distance from a to d = 7
Distance from a to e = 9
Distance from a to g = 18
Distance from a to f = 17

I'm so happy 😌

@edit: The code above is the version 1. Due community contributions (thanks!), it improved! Get the latest version here.

Posted on by:

Discussion

pic
Editor guide