Day14DockingData.kt

package adventofcode.year2020

import adventofcode.Puzzle
import adventofcode.PuzzleInput
import adventofcode.common.String.replaceAt

class Day14DockingData(customInput: PuzzleInput? = null) : Puzzle(customInput) {
    private val initializationInstructions by lazy { input.lines().map(InitializationInstruction::invoke) }

    override fun partOne() =
        initializationInstructions
            .fold(Pair("X".repeat(36), emptyMap<Long, Long>())) { (mask, memoryMap), instruction ->
                when (instruction) {
                    is MaskInstruction -> Pair(instruction.mask, memoryMap)
                    is MemoryInstruction ->
                        Pair(
                            mask,
                            memoryMap +
                                mapOf(
                                    instruction.address to
                                        (mask.indices)
                                            .fold(instruction.value.toString(2).padStart(mask.length, '0')) { result, index ->
                                                if (mask[index] == 'X') result else result.replaceAt(index, mask[index])
                                            }.toLong(2),
                                ),
                        )
                }
            }.second.values.sum()

    override fun partTwo() =
        initializationInstructions
            .fold(Pair("0".repeat(36), emptyMap<Long, Long>())) { (mask, memoryMap), instruction ->
                when (instruction) {
                    is MaskInstruction -> Pair(instruction.mask, memoryMap)
                    is MemoryInstruction -> {
                        val addressMask =
                            (mask.indices)
                                .fold(instruction.address.toString(2).padStart(mask.length, '0')) { result, index ->
                                    if (mask[index] == '0') result else result.replaceAt(index, mask[index])
                                }

                        Pair(
                            mask,
                            memoryMap +
                                (addressMask.indices)
                                    .fold(listOf(addressMask)) { acc, index ->
                                        when {
                                            addressMask[index] != 'X' -> acc
                                            else -> acc.flatMap { listOf(it.replaceAt(index, '0'), it.replaceAt(index, '1')) }
                                        }
                                    }
                                    .map { it.toLong(2) to instruction.value },
                        )
                    }
                }
            }.second.values.sum()

    companion object {
        private sealed class InitializationInstruction {
            companion object {
                operator fun invoke(input: String): InitializationInstruction {
                    val parts = input.split(" = ")

                    return when (parts.first()) {
                        "mask" -> MaskInstruction(parts.last())
                        else -> MemoryInstruction(parts.first().substring(4 until parts.first().length - 1).toLong(), parts.last().toLong())
                    }
                }
            }
        }

        private data class MaskInstruction(val mask: String) : InitializationInstruction()

        private data class MemoryInstruction(val address: Long, val value: Long) : InitializationInstruction()
    }
}