Dijkstra’s Algorithm in Kotlin

Dijkstra’s algorithm answers the question:

What is the shortest path from a start vertex to every other vertex in the graph?

This is useful—we need to do things like route messages through a phone network, and find the best route from point A to point B.

It hinges on the observation that any subpath of a shortest path must also be a shortest path (I encourage you to think about why). Here is my favorite explanation of it, and here is a useful visualization.

My code is more focused on education, mirroring the way I learned it originally at GT. To speed it up, you can make the operation of getting the next vertex to explore (v, in the code) done with a min heap.

Some other good implementations to compare with are Rosetta code, and the one in Kotlin Algorithm Club (what an awesome project!).

So, first let’s write a Graph class, so the algorithm has something to act upon.

data class Graph<T>(
    val vertices: Set<T>,
    val edges: Map<T, Set<T>>,
    val weights: Map<Pair<T, T>, Int>
)

Without going too much into the details of initialization and providing a simple API, this is essentially my Graph class (on Github I add an additional constructor).

Now, let’s actually take a look at the algorithm.

1
fun <T> dijkstra(graph: Graph<T>, start: T): Map<T, T?> {
2
    val S: MutableSet<T> = mutableSetOf() // a subset of vertices, for which we know the true distance
3

4
    val delta = graph.vertices.map { it to Int.MAX_VALUE }.toMap().toMutableMap()
5
    delta[start] = 0
6

7
    val previous: MutableMap<T, T?> = graph.vertices.map { it to null }.toMap().toMutableMap()
8

9
    while (S != graph.vertices) {
10
        val v: T = delta
11
            .filter { !S.contains(it.key) }
12
            .minBy { it.value }!!
13
            .key
14

15
        graph.edges.getValue(v).minus(S).forEach { neighbor ->
16
            val newPath = delta.getValue(v) + graph.weights.getValue(Pair(v, neighbor))
17

18
            if (newPath < delta.getValue(neighbor)) {
19
                delta[neighbor] = newPath
20
                previous[neighbor] = v
21
            }
22
        }
23

24
        S.add(v)
25
    }
26

27
    return previous.toMap()
28
}
29

30
fun <T> shortestPath(shortestPathTree: Map<T, T?>, start: T, end: T): List<T> {
31
    fun pathTo(start: T, end: T): List<T> {
32
        if (shortestPathTree[end] == null) return listOf(end)
33
        return listOf(pathTo(start, shortestPathTree[end]!!), listOf(end)).flatten()
34
    }
35

36
    return pathTo(start, end)
37
}

Man, Kotlin has fantastic API’s for functional programming.

Here’s the test:

1
@Test
2
fun shouldCalculateCorrectShortestPaths() {
3
    val weights = mapOf(
4
        Pair("A", "B") to 2,
5
        Pair("A", "C") to 8,
6
        Pair("A", "D") to 5,
7
        Pair("B", "C") to 1,
8
        Pair("C", "E") to 3,
9
        Pair("D", "E") to 2
10
    )
11

12
    val start = "A"
13
    val shortestPathTree = dijkstra(Graph(weights), start)
14

15
    Assert.assertEquals(listOf(start, "B", "C"), shortestPath(shortestPathTree, start, "C"))
16
    Assert.assertEquals(listOf(start, "B", "C", "E"), shortestPath(shortestPathTree, start, "E"))
17
    Assert.assertEquals(listOf(start, "D"), shortestPath(shortestPathTree, start, "D"))
18
}

If you’d like to run the code, clone the git repo and get started using gradle.