Day 11 Part 2

2024/src/main/kotlin/day11.kt

@@ -1,31 +1,28 @@
 fun main() {
     val input = readInputString("day11.txt").split(" ").map { it.toLong() }
 
-    val part1 = (1..25).fold(input) { result, _ ->
-        result.blink()
-    }.count()
-
+    val part1 = input.blink(25)
     println("Part 1: $part1")
 
-
-    val part2 = (1..75).fold(input) { result, _ ->
-        result.blink()
-    }.count()
-
+    val part2 = input.blink(75)
     println("Part 2: $part2")
-
 }
 
-fun List<Long>.blink() = flatMap {
-    val str = "$it"
+fun List<Long>.blink(count: Int) = fold(0L) { acc, stone -> acc + blink(stone, count) }
+
+val memo = mutableMapOf<Pair<Long, Int>, Long>()
+fun blink(stone: Long, count: Int): Long = memo.getOrPut(Pair(stone, count)) {
+    if (count == 0) return 1
+
+    val str = "$stone"
     val strLen = str.count()
 
     when {
-        it == 0L -> listOf(1L)
-        strLen % 2 == 0 -> listOf(
-            str.substring(0, strLen / 2).toLong(),
-            str.substring(strLen / 2).toLong()
-        )
-        else -> listOf(it * 2024)
+        stone == 0L -> blink(1L, count - 1)
+        strLen % 2 == 0 ->
+            blink(str.substring(0, strLen / 2).toLong(), count - 1) +
+                    blink(str.substring(strLen / 2).toLong(), count - 1)
+
+        else -> blink(stone * 2024, count - 1)
     }
 }