Collection processing in Kotlin: Grouping elements
// 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 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)]
}
Marcin Moskala is a highly experienced developer and Kotlin instructor as the founder of Kt. Academy, an official JetBrains partner specializing in Kotlin training, is known for his significant contributions to the Kotlin community. Moskala is the author of several widely recognized books, including "Effective Kotlin," "Kotlin Coroutines," "Functional Kotlin," "Advanced Kotlin," "Kotlin Essentials," and "Android Development with Kotlin."
Beyond his literary achievements, Moskala is the author of the largest Medium publication dedicated to Kotlin. As a respected speaker, he has been invited to share his insights at numerous programming conferences, including events such as Droidcon and the prestigious Kotlin Conf, the premier conference dedicated to the Kotlin programming language.
Owen has been developing software since the mid 1990s and remembers the productivity of languages such as Clipper and Borland Delphi.
Since 2001, He moved to Web, Server based Java and the Open Source revolution.
With many years of commercial Java experience, He picked up on Kotlin in early 2015.
After taking detours into Clojure and Scala, like Goldilocks, He thinks Kotlin is just right and tastes the best.
Owen enthusiastically helps Kotlin developers continue to succeed.
