Idiomatic Kotlin: Solving Advent of Code Puzzles, Set Operations

Idiomatic Kotlin: Solving Advent of Code Puzzles, Set Operations

Grouping and counting the characters in strings and collections is a very typical task given in coding interviews. Usually, the solutions for such tasks are quite simple. In this blog post, we learn about a few useful functions in Kotlin that make set operations really simple.

The task: custom customs

The task for the Advent of Code Day 6 challenge requires us to count the answers on customs declaration forms. We are provided with the results of the filled-in forms for different groups of people. Below is the original sample provided in the AoC Day 6 challenge:

abc

a
b
c

ab
ac

a
a
a
a

b

This list represents answers from 5 groups:

  1. The first group contains one person who answered “yes” to 3 questions: a, b, and c.
  2. The second group contains three people who collectively answered “yes” to 3 questions: a, b, and c.
  3. The third group contains two people who collectively answered “yes” to 3 questions: a, b, and c.
  4. The fourth group contains four people who collectively answered “yes” to only 1 question, a.
  5. The last group contains one person who answered “yes” to only 1 question, b.

Our task is to calculate the total number of “yes” answers in each group. For the example above, that would be 3 + 3 + 3 + 1 + 1 = 11.

Sounds easy, right? It seems like we could solve this task by just iterating the data, counting, and summing it up to get the final result. Let’s try that!

Solving the puzzle

First, we need to read the data. We can use the File.readText() function and then trim off the whitespace at the beginning and the end of the file. As per the description, group answers are always separated by blank lines, so we’ll use two line separators to split our input into separate groups:

val newLine = System.lineSeparator()

val groups: List<string> = File("src/day06/input.txt")
   .readText()
   .trim()
   .split("$newLine$newLine")</string>

The result is a list of strings. Each string in the list represents the data about a group of answers to the questions from the customs declaration form. 

Our goal is to count the unique answers within the group. We don’t need to count the answers for each individual person. This boils down to counting the unique characters in the string that represents the answers of each group. 

A naive approach to solving the task above would be to iterate over the list of strings, eliminate duplicate characters by converting each string to a set, and calculate the size of all the sets. The sum would be the final answer.

var total = 0
for (group in groups) {
   total += group.replace(newLine, "").toSet().size
}

The group of answers is one string, but it may include answers from multiple people, one line per person. Line breaks should not be taken into account when counting the unique characters. Therefore, I’m removing the line breaks before converting them to a set.

The solution is simple, but can it be improved in any way? As you might have guessed –  it can! In fact, summing up numbers by iterating over a data structure is a pretty common task. IntelliJ IDEA would even suggest using the sumOf function if we place the caret on the for loop and press Alt+Enter.

val total = groups.sumOf {
   it.replace(newLine, "").toSet().size
}

The sumOf function eliminates the need for the mutable variable to calculate the total. The function comes from Kotlin’s standard library, and if we look at the implementation, it does exactly what we did above:

public inline fun <T> Iterable<T>.sumOf(selector: (T) -> Int): Int {
   var sum: Int = 0.toInt()
   for (element in this) {
       sum += selector(element)
   }
   return sum
}

The selector parameter is the function that returns a number. In our task, that’s the size of the set with unique answers for the group.

The puzzle continues

There’s a second part to the task. Just like in real life – there’s a change in the requirements! According to the story, we have misread the rules for counting the answers in the group. Instead of finding the questions to which anyone answered “yes”, we need to find the questions to which everyone in the group answered “yes”. For the same input:

abc

a
b
c

ab
ac

a
a
a
a

b

The rules for counting the answers are the following:

  • In the first group, everyone (the 1 person) answered “yes” to all 3 questions: a, b, and c.
  • In the second group, there is no question to which everyone answered “yes”.
  • In the third group, there is 1 question to which everyone answered “yes”, a. Since some people did not answer “yes” to b or c, the answers to these questions don’t count.
  • In the fourth group, everyone answered yes to 1 question only, a.
  • In the fifth group, everyone (the 1 person) answered “yes” to 1 question, b.

The final answer is 3 + 0 + 1 + 1 + 1 = 6

Let’s make a plan for solving this task. We need to iterate over the list of answer groups (1). Each group is represented by a multi-line string, and each line in a string represents answers by an individual person. Hence, we need to split the string to analyze the individual answers (2). We need to find the characters that are present in each line in a group (3). Each line is “a set of answers” by an individual person. So, we have a number of sets and we need to find the elements that are included in all the sets. This means we need to find the intersection (4) of the sets and reduce (5) the sets to only those elements that belong to the intersection.

Luckily, intersect is a standard library function in Kotlin. Additionally, the collection of sets can be reduced to a single set with the reduce function.

var total = 0
for (group in groups) { // (1)
   val answers: List<String> = group.split(newLine) // (2)
   val answerSets: List<Set<Char>> = answers.map { it.toSet() } // (3)
   val intersection: Set<Char> = 
         answerSets.reduce /* (5) */ { a, b -> a intersect b /* (4) */}
   total += intersection.size
}

This is a reasonably straightforward approach for solving the task. However, we can definitely do better by using more functions from the Kotlin standard library.

We can logically split this code into two parts. First, we can shape the input into a data structure that can be used to implement the “business logic”. To do this, we need to split each string in the list that we have read from the file, and then we need to convert each element to a set. This can be implemented using the map function as follows:

val data = groups.map {
   it.split(newLine).map { s -> s.toSet() }
}

Next, we need to find the characters that are found in all the sets for each group. The reduce and intersect operations are still applicable for this part.

Let’s take a look at the intersect function first. Given two sets of characters, “abc” and “cde”, the intersection is only the character “c”:

val setA = "abc".toSet()
val setB = "cde".toSet()
val setC = setA intersect setB
println(setC) // prints [c]

If there are more sets, we can chain the operations by iterating the list of sets:

val setA = "abc".toSet()
val setB = "cde".toSet()
val setC = "cfg".toSet()
val setD = "czy".toSet()
val list = listOf(setA, setB, setC, setD)

var result = setA
for (set in list) {
    result = result intersect set
}
println(result) // prints [c]

The loop above is effectively the same as using the reduce function with the intersect operation as follows:

val setA = "abc".toSet()
val setB = "cde".toSet()
val setC = "cfg".toSet()
val setD = "czy".toSet()
val list = listOf(setA, setB, setC, setD)

val result = list.reduce { a, b -> a intersect b }

println(result) // prints [c]

All we need to do now is call the count() function on the result to get the number of answers that were the same in the specific group. The code looks like this:

val data = groups.map { 
   it.split(newLine).map(String::toSet)
}

var total = 0
for (listOfSets in data) {
   total += listOfSets.reduce { a, b -> a intersect b}.count()
}

We can see a familiar pattern for summing up the total. This is definitely a hint that we can use the sumOf function to calculate the same:

val total = data.sumOf { it ->
   it.reduce { a, b -> a intersect b }.count()
}

Now we can chain the operations to put all the calculations in one pipeline:

val total = groups.map {
   it.split(newLine).map(String::toSet)
}.sumOf { 
   it.reduce { a, b -> a intersect b }.count()
}

This code looks concise, but it is not perfectly readable yet. You can see that we are using the implicit variable name, 'it', in both code blocks. First, in the part where we transform the grouped answers into sets. Then, when working with the sets to find the intersections. The variable type in the different blocks is also different.

To make this more readable, we recommend giving an explicit name to the lambda argument, instead of 'it', if the code blocks are chained like in our example. 

val total = groups.map { group ->
   group.split(newLine).map(String::toSet)
}.sumOf { groupAnswers ->
   groupAnswers.reduce { a, b -> a intersect b }.count()
}

At this point, we can make an interesting observation: the first part of our task can be solved exactly like the second part, but using union instead of intersect. So here’s a crazy idea, what if we wanted to support the different operations within the same code. For this, we can extract the logic into a function and use the operation as a parameter for the function as follows:

private fun transformAndReduce(groups: List<String>, operation: (Set<Char>, Set<Char>) -> Set<Char>) =
   groups.map { lines ->
       lines.split(newLine).map(String::toSet)
   }.sumOf { characters ->
       characters.reduce { a, b -> operation(a , b) }.count()
   }

val firstAnswer = transformAndReduce(groups, Set<Char>::union)
val secondAnswer = transformAndReduce(groups, Set<Char>::intersect)

For this specific task, this might be overthinking things a bit. However, it demonstrates the use of functional types as parameters quite well.

Summary

Today’s puzzle did not require a complex solution. The final solution makes use of a few standard library functions: map, reduce, sumOf, intersect, and union. Using these functions we are able to elegantly compose a data manipulation pipeline to solve the puzzle. 

The scenario is quite common in real-life applications. Hopefully, the demo code will provide you with a good reference that you can use for your everyday tasks.

You can find the final solution in the GitHub repository here: https://github.com/kotlin-hands-on/advent-of-code-2020/blob/master/src/day06/day6.kt