Idiomatic Kotlin: Solving Advent of Code Puzzles, Traversing Trees

Pasha Finkelshteyn

Let’s describe our task briefly. We are airport luggage throwers handlers, and as often happens, there are some new regulations. This time, they stipulate that all bags be color-coded and contain a specific number of color-coded bags inside them. Weird, right? Well, rules are rules, so we must comply. Let’s see what they look like:

Our mission is to find how many bags eventually contain at least one shiny gold bag.

It would be hard to find potential container bags for our shiny gold bag by hand, but programming will help us! To solve this task, we will implement a couple of search algorithms.

Part 1

We will work with a tree of bags. Trees are a good way to represent structures where things are contained in other things. The most straightforward representation of the tree in our case will be a Map of Colors to Rules, where the key is a color or the contained bag and the rule is a set of colors of bags that potentially can contain our bag.

Let’s define a couple of type aliases for that.

typealias Color = String
typealias Rule = Set<String>

We also know that the most important bag for us is the shiny gold bag. Let’s put that into a constant.

private const val SHINY_GOLD = "shiny gold"

And now, let’s define the function that will return our tree of rules.

fun main() {
    val rules: Map<Color, Rule> = buildBagTree()
}

Now it’s time to start the first part of the solution: file parsing. Looking at our file again, we see each line is structured in the same way, including:

• The name of the container bag
• The word “bags”
• The word “contain”
• A single-digit number representing how many bags it contains
• A comma-separated list of bags that the container bag contains

To process all the lines, we’ll remove all of the words “bags” or “bag”, all of the dots, and then the word “contain” from our string. Now, everything to the left of “contain” will be the container bag’s name and color, and everything to the right of “contain” will be the name of contained bags. In this part of the puzzle, we don’t need the numbers that the rules mention, so let’s remove them too.

We’ll use regular expressions to remove these unnecessary parts. The regular expression for deleting single digits is simple, but the one for removing “bag” and “bags” with the following dot is a bit more complex:

private fun buildBagTree(): Map<Color, Rule> {
    File("src/day07/input.txt")
        .forEachLine { line ->
            line
                .replace(Regex("\\d+"), "")
                .replace(Regex("bags?\\.?"), "")
        }
    TODO()
}

After the pruning is complete, we need to split the string at “contain” and trim each resulting string.

private fun buildBagTree(): Map<Color, Rule> {
    File("src/day07/input.txt")
        .forEachLine { line ->
            line
                .replace(Regex("\\d+"), "")
                .replace(Regex("bags?\\.?"), "")
                .split("contain")
                .map { it.trim() }
        }
    TODO()
}

Let’s put the computation result into two different variables with a destructuring declaration:

private fun buildBagTree(): Map<Color, Rule> {
    File("src/day07/input.txt")
        .forEachLine { line ->
            val (parent, allChildren) = line
                .replace(Regex("\\d+"), "")
                .replace(Regex("bags?\\.?"), "")
                .split("contain")
                .map { it.trim() }
        }
    TODO()
}

The parent variable will contain the first element of the array, and allChildren will contain the second element of the array. Both will be strings.

Now, we need to convert the allChildren string into separate elements. We split it at the comma and trim again for each child.

private fun buildBagTree(): Map<Color, Rule> {
    File("src/day07/input.txt")
        .forEachLine { line ->
            val (parent, allChildren) = line
                .replace(Regex("\\d+"), "")
                .replace(Regex("bags?\\.?"), "")
                .split("contain")
                .map { it.trim() }
        }
            val childrenColors = allChildren.split(',').map { it.trim() }.toSet()
    TODO()
}

Our function needs to return a tree, so let’s define one:

private fun buildBagTree(): Map<Color, Rule> {
    val rules = hashMapOf<Color, Rule>()
    File("src/day07/input.txt")
        .forEachLine { line ->
            val (parent, allChildren) = line
                .replace(Regex("\\d+"), "")
                .replace(Regex("bags?\\.?"), "")
                .split("contain")
                .map { it.trim() }
            val childrenColors = allChildren.split(',').map { it.trim() }.toSet()
        }
    TODO()
}

Now, for each child found, we’ll check whether there is already a rule for it. If there is one, we add one more parent to it, and if not, we create a new rule. Remember that the rule is just a typealias to the set of strings.

A special function on maps called compute allows us to calculate the next value for any given key. Let’s make use of it. Note that this function is JVM-specific, so if you’re using a different platform you’ll need to use something else, like an if expression.

In our case, the key is the color of the child. The compute function accepts two arguments: the first is a key we are looking for, and the second is a lambda that, in turn, takes two arguments, our key and the current value.

private fun buildBagTree(): Map<Color, Rule> {
    val rules = hashMapOf<Color, Rule>()
    File("src/day07/input.txt")
        .forEachLine { line ->
            val (parent, allChildren) = line
                .replace(Regex("\\d+"), "")
                .replace(Regex("bags?\\.?"), "")
                .split("contain")
                .map { it.trim() }
            val childrenColors = allChildren.split(',').map { it.trim() }.toSet()
            for (childColor in childrenColors) {
                rules.compute(childColor) { key, value ->
                        TODO()
                }
            }
        }
    TODO()
}

The key type is String, and the value type is Set<String>?. Now we should define separate logics for the situation when the current value is null and when the current value is not null. When the rule is null, we create a new rule and return it as a new value. If the value is not null, we add a new parent to the existing rule.

private fun buildBagTree(): Map<Color, Rule> {
    val rules = hashMapOf<Color, Rule>()
    File("src/day07/input.txt")
        .forEachLine { line ->
            val (parent, allChildren) = line
                .replace(Regex("\\d+"), "")
                .replace(Regex("bags?\\.?"), "")
                .split("contain")
                .map { it.trim() }
            val childrenColors = allChildren.split(',').map { it.trim() }.toSet()
            for (childColor in childrenColors) {
                rules.compute(childColor) { key, value ->
                        if (current == null) setOf(parent)
                        else current + parent
                }
            }
        }
    TODO()
}

The rule is represented with an immutable set, so we can use the plus operator, which will create a new set for us. value and key aren’t very meaningful names, so let’s rename them. It also looks like the key is not used at all here because we already know our key. Let’s replace it with a special underscore value so it won’t be underlined as unused. Parsing is finished, so let’s make this function return our rules:

private fun buildBagTree(): Map<Color, Rule> {
    val rules = hashMapOf<Color, Rule>()
    File("src/day07/input.txt")
        .forEachLine { line ->
            val (parent, allChildren) = line
                .replace(Regex("\\d+"), "")
                .replace(Regex("bags?\\.?"), "")
                .split("contain")
                .map { it.trim() }
            val childrenColors = allChildren.split(',').map { it.trim() }.toSet()
            for (childColor in childrenColors) {
                rules.compute(childColor) { _, current ->
                    if (current == null) setOf(parent)
                    else current + parent
                }
            }
        }
    return rules
}

The next part is the actual search. How do we do that? Let’s use a depth-first search, a good algorithm for traversing trees or graphs.

For each iteration, we’ll define for which bags we need to find potential parents. The first step is to define our known bags, which is currently the “shiny gold” bag (result will be stored in the known variable). In the second step, we will find potential containers found on the first iteration (for the ”shiny gold” we know that there are at least some), et cetera, et cetera, until we finish the search. The result of an iteration will be called next.

If all of the next bags are already known, then we’ll break the cycle.

Otherwise, we’ll redefine our already known bags and the next step with new values. The most exciting part here is defining the new step. It will be the potential containers of each bag in the current toFind set. For each item in the toFind set, we’ll find its containers in rules, and then flatten the search results and convert them to a set to remove any duplicates.

fun findContainersDFS(rules: Map<Color, Rule>): Set<Color> {
    var known = setOf(SHINY_GOLD)
    var next = setOf(SHINY_GOLD) + rules[SHINY_GOLD]!!
    while (true) {
            val toFind = next - known
            if (toFind.isEmpty()) break
            known = known + next
            next = toFind.mapNotNull { rules[it] }.flatten().toSet()
    }
    return known - SHINY_GOLD
}

It looks like our iterative search is finished. Let’s return the results from our search function and output the final result – the known bags without the shiny gold bag itself.

fun main() {
    val rules: Map<Color, Rule> = buildBagTree()
    val containers = findContainersDFS(rules)
    println(containers)
    println()
    println(containers.size)
}

I run the program and my output is 274. It’s the correct answer for how many bags contain at least one shiny gold bag.

Now let’s move on to the second part of the puzzle.

Part 2

This time the input (the list of rules) is the same, but we need to count how many bags our shiny gold bag will eventually contain.

We’ll perform the same calculation, but this time we’ll go down the tree and use the numbers (which we threw out last time) in order to count how many bags each bag will contain, and we’ll do this recursively.

Let’s look at our input again. Rereading the rules gives us new information: some bags won’t contain other bags, according to lines such as “shiny green bags contain no other bags“. That would be the end of the recursion for that kind of bag.

Now we’re going to build the same tree, but the color of the container will be the key, and the info about children will be the rule. Let’s implement the parsing logic:

private fun buildBagTree(): Map<Color, Rule> {
    val rules = hashMapOf<Color, Rule>()
    File("src/day07/input.txt")
        .forEachLine { line ->
            val (parent, allChildren) = line.replace(Regex("bags?\\.?"), "")
                .split("contain")
                .map { it.trim() }

        }
    return rules
}

Now let’s compute our rule for the current parent. If there is a phrase “no other” in the “allChildren” variable, the rule will be empty. Otherwise, it will contain children split by comma, trimmed, and converted into a set.

private fun buildBagTree(): Map<Color, Rule> {
    val rules = hashMapOf<Color, Rule>()
    File("src/day07/input.txt")
        .forEachLine { line ->
            val (parent, allChildren) = line.replace(Regex("bags?\\.?"), "")
                .split("contain")
                .map { it.trim() }
            val rule =
                if (allChildren.contains("no other")) emptySet()
                else allChildren.split(',').map { it.trim() }.toSet()
            rules[parent] = rule
        }
    return rules
}

Now we need to count the children inside the shiny gold bag with the help of a breadth-first search. A breadth-first search is an algorithm for searching a tree data structure for a node that satisfies a given property. It starts at the tree root and explores all nodes at the present depth before moving on to the nodes of the next level.

Let’s create a function for it – an extension function for our map for the sake of readability. As a parameter, it will accept the color of the bag children which we are counting.

private fun Map<Color, Rule>.getChildrenCountBFS(color: Color): Int {
    TODO()
}

Remember how we kept the number of every child in the bag? Now let’s create a regular expression, which we’ll use to find the numbers in each child.

val digits = "\\d+".toRegex()

We extract all children from the current bag at each level of iteration. If there are no other bags in it, we return zero.

private fun Map<Color, Rule>.getChildrenCountBFS(color: Color): Int {
    val children = getOrDefault(color, setOf())
    if (children.isEmpty()) return 0
      TODO()
}

Now we declare the counter and start counting bags inside each child. First, we extract the count of this child in the parent bag to the variable “count”.

Second, we replace this count with nothing in the child’s name and trim the resulting name. We add to the total count and multiply by the number of bags inside each bag. On each level of recursion, we return the total after the loop:

private fun Map<Color, Rule>.getChildrenCountBFS(color: Color): Int {
    val children = getOrDefault(color, setOf())
    if (children.isEmpty()) return 0
    var total = 0
    for (child in children) {
            val count = digits.findAll(child).first().value.toInt()
            val bag = digits.replace(child, "").trim()
            total += count + count * getChildrenCountBFS(bag)
    }
    return total
}

That’s it: the result is 158,730 and it’s correct!

This is what we learned about today:

• How to use regular expressions
• JVM-specific compute methods that allow us to dynamically calculate the value in the map based on the current value
• The forEachLine method
• Two types of search algorithms:
  • Depth-first search
  • Breadth-first search

Subscribe to Kotlin Blog updates

Subscribe form

By submitting this form, I agree to the JetBrains Privacy Policy Notification icon

By submitting this form, I agree that JetBrains s.r.o. ("JetBrains") may use my name, email address, and location data to send me newsletters, including commercial communications, and to process my personal data for this purpose. I agree that JetBrains may process said data using third-party services for this purpose in accordance with the JetBrains Privacy Policy. I understand that I can revoke this consent at any time in my profile. In addition, an unsubscribe link is included in each email.

Submit

Thanks, we've got you!