day(15): move Dijkstra to separate class

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

src/year2021/day15/Day15.kt

@@ -3,39 +3,52 @@ 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>> =
+open class Dijkstra(
-    (y0 - 1..y0 + 1)
+    private val input: List<List<Int>>
-        .flatMap { y -> (x0 - 1..x0 + 1).map { x -> Pair(y, x) } }
+) {
-        .filter { (y, x) ->
+    private val queue: PriorityQueue<Pair<Int, Coordinate>> =
-            (y != y0 || x != x0) && (y == y0 || x == x0) && (y >= 0 && y < 5 * input.size) && (x >= 0 && x < 5 * input[0].size)
+        PriorityQueue<Pair<Int, Coordinate>> { a, b ->
-        }
-
-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>()
+    private val distances = mutableMapOf<Coordinate, Int>()
-    val processed = mutableSetOf<Coordinate>()
+    private val processed = mutableSetOf<Coordinate>()
 
-    queue.add(Pair(0, Coordinate(0, 0)))
+    fun getDistance(coordinate: Coordinate): Int =
-    distances[Coordinate(0, 0)] = 0
+        distances[coordinate] ?: Int.MAX_VALUE
-    for (y in input.indices) {
+
-        for (x in input[y].indices) {
+    open fun getCost(coordinate: Coordinate): Int =
-            queue.add(Pair(Int.MAX_VALUE, Coordinate(y, x)))
+        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
     }
 
+    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)) {
@@ -43,75 +56,49 @@ fun part1(input: List<List<Int>>): Int {
             }
             processed.add(coordinate)
 
-        if (coordinate == final) {
+            if (coordinate == end) {
-            return distance
+                return this
             }
 
-        for ((y, x) in validAdjacentIndices(input, coordinate.y, coordinate.x)) {
+            neighbours(coordinate).forEach { relax(distance, it) }
-            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 {
+        return this
-    val riskLevel = input[y % input.size][x % input[0].size]
+    }
+}
+
+fun part1(input: List<List<Int>>): Int =
+    Coordinate(input.size - 1, input.last().size - 1).let { end ->
+        Dijkstra(input)
+            .run(
+                Coordinate(0, 0),
+                end
+            ).getDistance(end)
+    }
+
+class DijkstraOnExtended(private val input: List<List<Int>>) : Dijkstra(input) {
+    override fun neighbours(coordinate: Coordinate): Iterable<Coordinate> =
+        coordinate.adjacent().filter { (y, x) ->
+            (y >= 0 && y < 5 * input.size) && (x >= 0 && x < 5 * input[0].size)
+        }
+
+    override fun getCost(coordinate: Coordinate): Int {
+        val (y, x) = coordinate
+
         val added = y / input.size + x / input[0].size
+        val riskLevel = input[y % input.size][x % input[0].size] + added
 
-    if (riskLevel + added > 9) {
+        return if (riskLevel > 9) riskLevel - 9 else riskLevel
-        return riskLevel + added - 9
-    }
-
-    return riskLevel + added
-}
-
-fun part2(input: List<List<Int>>): Int {
-    val final = Coordinate(5 * input.size - 1, 5 * 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) {
-            for (i in 0 until 5) {
-                queue.add(Pair(Int.MAX_VALUE, Coordinate(i * y, i * x)))
-            }
     }
 }
 
-    while (queue.isNotEmpty()) {
+fun part2(input: List<List<Int>>): Int =
-        val (distance, coordinate) = queue.remove()!!
+    Coordinate(5 * input.size - 1, 5 * input.last().size - 1).let { end ->
-        if (processed.contains(coordinate)) {
+        DijkstraOnExtended(input)
-            continue
+            .run(
-        }
+                Coordinate(0, 0),
-        processed.add(coordinate)
+                end
-
+            ).getDistance(end)
-        if (coordinate == final) {
-            return distance
-        }
-
-        for ((y, x) in validAdjacentIndicesExtended(input, coordinate.y, coordinate.x)) {
-            val neighbour = Coordinate(y, x)
-            val alternative = distance + getDistance(input, y, x)
-
-            if (alternative < (distances[neighbour] ?: Int.MAX_VALUE)) {
-                distances[neighbour] = alternative
-                queue.add(Pair(alternative, neighbour))
-            }
-        }
-    }
-    return Int.MAX_VALUE
     }
 
 fun main() {