day(15): move Dijkstra to separate class

Signed-off-by: Matej Focko <mfocko@redhat.com>

src/year2021/day15/Day15.kt

@@ -3,117 +3,104 @@ package year2021.day15
 import readInput
 import java.util.*
 
-fun validAdjacentIndices(input: List<List<Int>>, y0: Int, x0: Int): Iterable<Pair<Int, Int>> =
+data class Coordinate(val y: Int, val x: Int) {
-    (y0 - 1..y0 + 1)
+    fun adjacent(): Iterable<Coordinate> =
-        .flatMap { y -> (x0 - 1..x0 + 1).map { x -> Pair(y, x) } }
+        listOf(
-        .filter { (y, x) ->
+            Coordinate(y - 1, x),
-            (y != y0 || x != x0) && (y == y0 || x == x0) && (y >= 0 && y < input.size) && (x >= 0 && x < input[y].size)
+            Coordinate(y + 1, x),
-        }
+            Coordinate(y, x + 1),
-
+            Coordinate(y, x - 1)
-fun validAdjacentIndicesExtended(input: List<List<Int>>, y0: Int, x0: Int): Iterable<Pair<Int, Int>> =
+        )
-    (y0 - 1..y0 + 1)
-        .flatMap { y -> (x0 - 1..x0 + 1).map { x -> Pair(y, x) } }
-        .filter { (y, x) ->
-            (y != y0 || x != x0) && (y == y0 || x == x0) && (y >= 0 && y < 5 * input.size) && (x >= 0 && x < 5 * input[0].size)
-        }
-
-data class Coordinate(val y: Int, val x: Int)
-
-fun part1(input: List<List<Int>>): Int {
-    val final = Coordinate(input.size - 1, input.last().size - 1)
-
-    val queue = PriorityQueue<Pair<Int, Coordinate>>() { a, b ->
-        a.first - b.first
-    }
-    val distances = mutableMapOf<Coordinate, Int>()
-    val processed = mutableSetOf<Coordinate>()
-
-    queue.add(Pair(0, Coordinate(0, 0)))
-    distances[Coordinate(0, 0)] = 0
-    for (y in input.indices) {
-        for (x in input[y].indices) {
-            queue.add(Pair(Int.MAX_VALUE, Coordinate(y, x)))
-        }
-    }
-
-    while (queue.isNotEmpty()) {
-        val (distance, coordinate) = queue.remove()!!
-        if (processed.contains(coordinate)) {
-            continue
-        }
-        processed.add(coordinate)
-
-        if (coordinate == final) {
-            return distance
-        }
-
-        for ((y, x) in validAdjacentIndices(input, coordinate.y, coordinate.x)) {
-            val neighbour = Coordinate(y, x)
-            val alternative = distance + input[y][x]
-
-            if (alternative < (distances[neighbour] ?: Int.MAX_VALUE)) {
-                distances[neighbour] = alternative
-                queue.add(Pair(alternative, neighbour))
-            }
-        }
-    }
-    return Int.MAX_VALUE
 }
 
-fun getDistance(input: List<List<Int>>, y: Int, x: Int): Int {
+open class Dijkstra(
-    val riskLevel = input[y % input.size][x % input[0].size]
+    private val input: List<List<Int>>
-    val added = y / input.size + x / input[0].size
+) {
+    private val queue: PriorityQueue<Pair<Int, Coordinate>> =
+        PriorityQueue<Pair<Int, Coordinate>> { a, b ->
+            a.first - b.first
+        }
+    private val distances = mutableMapOf<Coordinate, Int>()
+    private val processed = mutableSetOf<Coordinate>()
 
-    if (riskLevel + added > 9) {
+    fun getDistance(coordinate: Coordinate): Int =
-        return riskLevel + added - 9
+        distances[coordinate] ?: Int.MAX_VALUE
+
+    open fun getCost(coordinate: Coordinate): Int =
+        input[coordinate.y][coordinate.x]
+
+    private fun relax(distance: Int, neighbour: Coordinate): Boolean {
+        val alternative = distance + getCost(neighbour)
+        if (alternative >= getDistance(neighbour)) {
+            return false
+        }
+
+        distances[neighbour] = alternative
+        queue.add(Pair(alternative, neighbour))
+        return true
     }
 
-    return riskLevel + added
+    open fun neighbours(coordinate: Coordinate): Iterable<Coordinate> =
+        coordinate.adjacent().filter { (y, x) ->
+            (y >= 0 && y < input.size) && (x >= 0 && x < input[y].size)
+        }
+
+    fun run(start: Coordinate, end: Coordinate): Dijkstra {
+        queue.add(Pair(0, start))
+        distances[start] = 0
+
+        while (queue.isNotEmpty()) {
+            val (distance, coordinate) = queue.remove()!!
+            if (processed.contains(coordinate)) {
+                continue
+            }
+            processed.add(coordinate)
+
+            if (coordinate == end) {
+                return this
+            }
+
+            neighbours(coordinate).forEach { relax(distance, it) }
+        }
+
+        return this
+    }
 }
 
-fun part2(input: List<List<Int>>): Int {
+fun part1(input: List<List<Int>>): Int =
-    val final = Coordinate(5 * input.size - 1, 5 * input.last().size - 1)
+    Coordinate(input.size - 1, input.last().size - 1).let { end ->
-
+        Dijkstra(input)
-    val queue = PriorityQueue<Pair<Int, Coordinate>>() { a, b ->
+            .run(
-        a.first - b.first
+                Coordinate(0, 0),
-    }
+                end
-    val distances = mutableMapOf<Coordinate, Int>()
+            ).getDistance(end)
-    val processed = mutableSetOf<Coordinate>()
-
-    queue.add(Pair(0, Coordinate(0, 0)))
-    distances[Coordinate(0, 0)] = 0
-    for (y in input.indices) {
-        for (x in input[y].indices) {
-            for (i in 0 until 5) {
-                queue.add(Pair(Int.MAX_VALUE, Coordinate(i * y, i * x)))
-            }
-        }
     }
 
-    while (queue.isNotEmpty()) {
+class DijkstraOnExtended(private val input: List<List<Int>>) : Dijkstra(input) {
-        val (distance, coordinate) = queue.remove()!!
+    override fun neighbours(coordinate: Coordinate): Iterable<Coordinate> =
-        if (processed.contains(coordinate)) {
+        coordinate.adjacent().filter { (y, x) ->
-            continue
+            (y >= 0 && y < 5 * input.size) && (x >= 0 && x < 5 * input[0].size)
-        }
-        processed.add(coordinate)
-
-        if (coordinate == final) {
-            return distance
         }
 
-        for ((y, x) in validAdjacentIndicesExtended(input, coordinate.y, coordinate.x)) {
+    override fun getCost(coordinate: Coordinate): Int {
-            val neighbour = Coordinate(y, x)
+        val (y, x) = coordinate
-            val alternative = distance + getDistance(input, y, x)
 
-            if (alternative < (distances[neighbour] ?: Int.MAX_VALUE)) {
+        val added = y / input.size + x / input[0].size
-                distances[neighbour] = alternative
+        val riskLevel = input[y % input.size][x % input[0].size] + added
-                queue.add(Pair(alternative, neighbour))
+
-            }
+        return if (riskLevel > 9) riskLevel - 9 else riskLevel
-        }
     }
-    return Int.MAX_VALUE
 }
 
+fun part2(input: List<List<Int>>): Int =
+    Coordinate(5 * input.size - 1, 5 * input.last().size - 1).let { end ->
+        DijkstraOnExtended(input)
+            .run(
+                Coordinate(0, 0),
+                end
+            ).getDistance(end)
+    }
+
 fun main() {
     val sample = readInput(15, "sample").map { row -> row.map { it.digitToInt() } }
     val input = readInput(15, "input").map { row -> row.map { it.digitToInt() } }