article banner (priority)

Collection processing in Kotlin: Grouping elements

This is a chapter from the book Functional Kotlin. You can find it on LeanPub.

// `partition` implementation from Kotlin stdlib inline fun <T> Iterable<T>.partition( predicate: (T) -> Boolean ): Pair<List<T>, List<T>> { val first = ArrayList<T>() val second = ArrayList<T>() for (element in this) { if (predicate(element)) { first.add(element) } else { second.add(element) } } return Pair(first, second) }

We have learned about the filter function, which returns a list of elements that satisfy a predicate, but what if we are interested in the elements that satisfy it as well as those that do not? In such a case, we use the partition method, which returns a pair of lists. The first list contains all elements that satisfy its predicate, and the second list contains those that do not. This pair can be then destructured into separate collections.

fun main() { val nums = listOf(1, 2, 6, 11) val partitioned: Pair<List<Int>, List<Int>> = nums.partition { it in 2..10 } println(partitioned) // ([2, 6], [1, 11]) val (inRange, notInRange) = partitioned println(inRange) // [2, 6] println(notInRange) // [1, 11] }

fun main() { val nums = (1..10).toList() val (smaller, bigger) = nums.partition { it <= 5 } println(smaller) // [1, 2, 3, 4, 5] println(bigger) // [6, 7, 8, 9, 10] val (even, odd) = nums.partition { it % 2 == 0 } println(even) // [2, 4, 6, 8, 10] println(odd) // [1, 3, 5, 7, 9] data class Student(val name: String, val passing: Boolean) val students = listOf( Student("Alex", true), Student("Ben", false), ) val (passing, failed) = students.partition { it.passing } println(passing) // [Student(name=Alex, passing=true)] println(failed) // [Student(name=Ben, passing=false)] }

groupBy

// `groupBy` implementation from Kotlin stdlib inline fun <T, K> Iterable<T>.groupBy( keySelector: (T) -> K ): Map<K, List<T>> { val destination = LinkedHashMap<K, MutableList<T>>() for (element in this) { val key = keySelector(element) val list = destination.getOrPut(key) { ArrayList<T>() } list.add(element) } return destination }

After presenting partition, I am often asked what we can do if we want to divide our collection into more than two groups. In such situations, we use groupBy, which groups elements by keys and returns a map from each key into a list of elements with this key (Map<K, List<E>>).

fun main() { val names = listOf("Marcin", "Maja", "Cookie") val byCapital = names.groupBy { it.first() } println(byCapital) // {M=[Marcin, Maja], C=[Cookie]} val byLength = names.groupBy { it.length } println(byLength) // {6=[Marcin, Cookie], 4=[Maja]} }

From my experience, when my colleagues ask me for help with more complex collection processing, pretty often what they are missing is groupBy. Here are a few tasks that require this operation1:

  • Count the number of users in each city, based on a list of users.
  • Find the number of points received by each team, based on a list of players.
  • Find the best option in each category, based on a list of options.
// Count the number of users in each city val usersCount: Map<City, Int> = users .groupBy { it.city } .mapValues { (_, users) -> users.size } // Find the number of points received by each team val pointsPerTeam: Map<Team, Int> = players .groupBy { it.team } .mapValues { (_, players) -> players.sumOf { it.points } } // Find the best resolution in each category val bestResolutionPerQuality: Map<Quality, Resolution> = formats.groupBy { it.quality } .mapValues { (_, formats) -> formats.maxOf { it.resolution } }

There is also the groupingBy method, which can be used as an alternative to groupBy. groupingBy is more efficient but also harder to use2.

You can reverse the groupBy method using flatMap. If you first use groupBy and then flatMap the values, you will have the same elements you started with (but possibly in a different order).

data class Player(val name: String, val team: String) fun main() { val players = listOf( Player("Alex", "A"), Player("Ben", "B"), Player("Cal", "A"), ) val grouped = players.groupBy { it.team } println(grouped) // {A=[Player(name=Alex, team=A), // Player(name=Cal, team=A)], // B=[Player(name=Ben, team=B)]} println(grouped.flatMap { it.value }) // [Player(name=Alex, team=A), Player(name=Cal, team=A), // Player(name=Ben, team=B)] }
1:

The mapValues function is a function on Map that transforms all values according to the transformation function.

2:

I described using groupingBy in Effective Kotlin, Item 53: Consider using groupingBy instead of groupBy.