From 4b84ce2803e057ce789bf29b2ae147a3b50dfc79 Mon Sep 17 00:00:00 2001
From: Chris Hodges <chrishod@axis.com>
Date: Tue, 10 Dec 2024 16:11:38 +0100
Subject: [PATCH] Optimized day 9 for no reason, part 2 now runs in about 50ms
 for really_evil_input.txt.

---
 src/aoc2024/Day09.kt | 40 ++++++++++++++--------------------------
 1 file changed, 14 insertions(+), 26 deletions(-)

diff --git a/src/aoc2024/Day09.kt b/src/aoc2024/Day09.kt
index ada61cc..d01488d 100644
--- a/src/aoc2024/Day09.kt
+++ b/src/aoc2024/Day09.kt
@@ -2,6 +2,7 @@ package aoc2024
 
 import println
 import readInput
+import java.util.*
 
 /*
 --- Day 9: Disk Fragmenter ---
@@ -45,24 +46,19 @@ fun main() {
     }
 
     fun part2(input: List<String>): Long {
-        val diskSize = input[0].sumOf { it - '0' }
-        val bitmap = IntArray(diskSize) { 0 }
         var pos = 0
         var isFile = true
         var fileId = 0
-        val freeList = HashMap<Int, MutableList<Int>>()
+        val freeList = Array<MutableSet<Int>>(10) { TreeSet() }
         val fileList = ArrayList<Pair<Int, Pair<Int, Int>>>()
         val newFileList = ArrayList<Pair<Int, Pair<Int, Int>>>()
         for (c in input[0]) {
             val size = c - '0'
             if (isFile) {
-                if (size == 0) println("Narf!")
                 fileList.add(fileId to (pos to size))
                 fileId++
             } else {
-                if (size > 0) {
-                    freeList.getOrPut(size) { ArrayDeque(0) }.add(pos)
-                }
+                if (size > 0) freeList[size].add(pos)
             }
             pos += size
             isFile = !isFile
@@ -70,35 +66,27 @@ fun main() {
 
         for (file in fileList.reversed()) {
             val (filePos, size) = file.second
-            var bestList: MutableList<Int>? = null
             var bestListSize = 0
             for (s in size..9) {
                 val list = freeList[s]
-                if (list?.isNotEmpty() == true && list[0] < filePos) {
-                    if (bestList == null || list[0] < bestList[0]) {
-                        bestList = list
-                        bestListSize = s
-                    }
+                val first = list.firstOrNull()
+                if ((first != null && first < filePos) &&
+                    (bestListSize == 0 || first < freeList[bestListSize].first())
+                ) {
+                    bestListSize = s
                 }
             }
-            if (bestList != null) {
-                val newPos = bestList.removeAt(0)
-                if (bestListSize > size) {
-                    val remBucket = freeList.getOrPut(bestListSize - size) { ArrayDeque(0) }
-                    remBucket.add(newPos + size)
-                    remBucket.sort()
-                }
+            if (bestListSize != 0) {
+                val newPos = freeList[bestListSize].first()
+                freeList[bestListSize].remove(newPos)
+                if (bestListSize > size) freeList[bestListSize - size].add(newPos + size)
                 newFileList.add(file.first to (newPos to size))
             } else {
                 newFileList.add(file)
             }
         }
-        for (file in newFileList) {
-            for (i in 0 until file.second.second) {
-                bitmap[i + file.second.first] = file.first
-            }
-        }
-        return bitmap.map(Int::toLong).reduceIndexed { index, acc, i -> acc + index * i }
+        // Gauß to the rescue!
+        return newFileList.sumOf { it.first * (it.second.first.toLong() * it.second.second + (it.second.second * (it.second.second - 1) / 2)) }
     }
 
     // test if implementation meets criteria from the description, like: