Advent of code 2021, day 12

aoc2021day12.kt

package net.chmielowski.aoc.day12

import net.chmielowski.aoc.day12.Vertex.*
import java.io.File

fun main() {
    val file = File("input.txt")
    val lines = file.readLines()

    println("Part 1: " + findPaths(lines, strategy = MaxOnce).size)
    println("Part 2: " + findPaths(lines, strategy = MaxOnceOrTwice).size)
}

fun findPaths(input: List<String>, strategy: VisitingSmallCaveStrategy): Set<Path> {
    val connections = parseConnections(input)
    return findPaths(connections = connections, strategy = strategy)
}

private fun parseConnections(input: List<String>) = input
    .flatMap(::parseLine)
    .groupBy { (from, _) -> from }
    .mapValues { (_, connections) ->
        connections
            .map { (_, to) -> to }
            .toSet()
    }
    .let(::Connections)

private fun parseLine(line: String): Set<Pair<Vertex, Vertex>> {
    val (rawA, rawB) = line.split("-")
    val a = parseVertex(rawA)
    val b = parseVertex(rawB)
    return setOf(a to b, b to a)
}

private fun parseVertex(raw: String) = when (raw) {
    "start" -> Start
    "end" -> End
    else -> when {
        raw.all(Char::isUpperCase) -> Cave.Large(raw)
        raw.all(Char::isLowerCase) -> Cave.Small(raw)
        else -> error("Invalid vertex $raw")
    }
}

private fun findPaths(connections: Connections, path: Path = Path(), strategy: VisitingSmallCaveStrategy): Set<Path> =
    if (path.isComplete) {
        setOf(path)
    } else {
        connections
            .from(path.last)
            .filter { vertex -> path.canVisit(vertex, strategy) }
            .flatMap { destination ->
                findPaths(connections, path + destination, strategy)
            }
            .toSet()
    }

sealed interface Vertex {

    object Start : Vertex {

        override fun toString() = "start"
    }

    sealed interface Cave : Vertex {

        data class Small(val name: String) : Cave {

            override fun toString() = name
        }

        data class Large(val name: String) : Cave {

            override fun toString() = name
        }
    }

    object End : Vertex {

        override fun toString() = "end"
    }
}

data class Path(private val vertices: List<Vertex> = listOf(Start)) {

    val last get() = vertices.last()

    val isComplete get() = last is End

    operator fun plus(next: Vertex) = copy(vertices = vertices + next)

    fun canVisit(vertex: Vertex, strategy: VisitingSmallCaveStrategy) = when (vertex) {
        Start -> false
        is Cave.Small -> strategy.canVisit(vertices.filterIsInstance<Cave.Small>(), vertex)
        is Cave.Large, End -> true
    }

    override fun toString() = vertices.joinToString(",")
}

data class Connections(private val map: Map<Vertex, Set<Vertex>>) {

    fun from(vertex: Vertex) = map[vertex] ?: error("No connection from $vertex")
}

fun interface VisitingSmallCaveStrategy {

    fun canVisit(previous: List<Cave.Small>, next: Cave.Small): Boolean
}

val MaxOnce = VisitingSmallCaveStrategy { previous, next ->
    !previous.contains(next)
}

val MaxOnceOrTwice = VisitingSmallCaveStrategy { previous, next ->
    val counts = previous
        .groupingBy { it }
        .eachCount()
    when (counts.getOrDefault(next, defaultValue = 0)) {
        2 -> false
        1 -> !counts.any { (_, count) -> count == 2 }
        0 -> true
        else -> error("Invalid path: $previous, adding $next")
    }
}