Dijkstra’s Algorithm in Kotlin

Apr 27, 2019

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.

Wow! You read the whole thing. People who make it this far sometimes want to receive emails when I post something new.

I also have an RSS feed.