1 files changed, 22 insertions, 44 deletions
diff --git a/src/Day06.kt b/src/Day06.kt
index d192c72..57e56f2 100644
--- a/src/Day06.kt
+++ b/src/Day06.kt
@@ -1,52 +1,30 @@
-fun main() {
-    fun part1(input: String): Int {
-        var fish = input.trim().split(",").map(String::toInt)
-
-        repeat(80) {
-            val pendingFish = mutableListOf<Int>()
-
-            fish = fish.map {
-                when (it) {
-                    0 -> {
-                        pendingFish.add(8)
-                        6
-                    }
-                    else -> {
-                        it - 1
-                    }
-                }
-            } + pendingFish
+fun calculateFishPopulation(input: String, days: Int): Long {
+    val fishCounts =
+        input
+            .trim()
+            .split(",")
+            .map(String::toInt)
+            .groupingBy { it }
+            .eachCount()
+            .mapValues { (_, v) -> v.toLong() }
+            .toMutableMap()
+
+    repeat(days) {
+        val readyToBirth = fishCounts.getOrDefault(0, 0)
+        repeat(8) {
+            fishCounts[it] = fishCounts.getOrDefault(it + 1, 0)
         }
-
-        return fish.size
+        fishCounts.merge(6, readyToBirth, Long::plus)
+        fishCounts[8] = readyToBirth
     }
 
-    fun part2(input: String): Long {
-        val fishCounts =
-            input
-                .trim()
-                .split(",")
-                .map(String::toInt)
-                .groupingBy { it }
-                .eachCount()
-                .mapValues { (_, v) -> v.toLong() }
-                .toMutableMap()
+    return fishCounts.values.sum()
+}
 
-        repeat(256) {
-            val readyToBirth = fishCounts.getOrDefault(0, 0)
-            fishCounts[0] = fishCounts.getOrDefault(1, 0)
-            fishCounts[1] = fishCounts.getOrDefault(2, 0)
-            fishCounts[2] = fishCounts.getOrDefault(3, 0)
-            fishCounts[3] = fishCounts.getOrDefault(4, 0)
-            fishCounts[4] = fishCounts.getOrDefault(5, 0)
-            fishCounts[5] = fishCounts.getOrDefault(6, 0)
-            fishCounts[6] = fishCounts.getOrDefault(7, 0) + readyToBirth
-            fishCounts[7] = fishCounts.getOrDefault(8, 0)
-            fishCounts[8] = readyToBirth
-        }
+fun main() {
+    fun part1(input: String): Int = calculateFishPopulation(input, 80).toInt()
 
-        return fishCounts.values.sum()
-    }
+    fun part2(input: String): Long = calculateFishPopulation(input, 256)
 
     val testInput = readInputAsString("Day06_test")
     check(part1(testInput) == 5934)