Day11DumboOctopus.kt

package adventofcode.year2021

import adventofcode.Puzzle
import adventofcode.PuzzleInput
import adventofcode.common.neighbors

class Day11DumboOctopus(customInput: PuzzleInput? = null) : Puzzle(customInput) {
    private val grid by lazy { input.lines().map { it.map { it.toString().toInt() } } }

    private fun List<List<Int>>.incrementAll() = map { it.map(Int::inc) }

    private fun List<List<Int>>.increment(points: List<Pair<Int, Int>>) =
        mapIndexed { y, row -> row.mapIndexed { x, col -> col + points.filter { it.first == x && it.second == y }.size } }

    private fun List<List<Int>>.flashes() = flatMapIndexed { y, row -> row.mapIndexedNotNull { x, col -> if (col > 9) x to y else null } }

    private fun List<List<Int>>.reset() = map { row -> row.map { col -> if (col > 9) 0 else col } }

    private fun List<List<Int>>.simulate(steps: Int) =
        generateSequence(0 to this) { (previousFlashes, previousGrid) ->
            val incrementedGrid = previousGrid.incrementAll()

            val (newFlashes, newGrid, _) =
                generateSequence(
                    Triple(0, incrementedGrid, emptySet<Pair<Int, Int>>()),
                ) { (currentFlashes, currentGrid, alreadyFlashed) ->
                    val flashes = currentGrid.flashes().minus(alreadyFlashed)
                    val neighbors = flashes.flatMap { currentGrid.neighbors(it.first, it.second, true) }

                    when (flashes.size) {
                        0 -> null
                        else -> Triple(currentFlashes + flashes.size, currentGrid.increment(neighbors), alreadyFlashed.plus(flashes))
                    }
                }
                    .last()

            previousFlashes + newFlashes to newGrid.reset()
        }
            .take(steps + 1)
            .last()

    override fun partOne() = grid.simulate(100).first

    override fun partTwo() =
        generateSequence(0 to grid) { (previousStep, previousGrid) ->
            previousStep + 1 to previousGrid.simulate(1).second
        }
            .first { (_, grid) -> grid.flatten().all { it == 0 } }
            .first
}