DEV Community

Cover image for [Tiny] Partition map in Kotlin
Petr Filaretov
Petr Filaretov

Posted on

[Tiny] Partition map in Kotlin

As you may know, Kotlin has a lot of useful functions to work with collections and maps. One such function is Iterable.partition() which splits a given iterable into a pair of lists by predicate. For instance:

val list = listOf(100, -5, 42, 25)
val (lessThan42, greaterThanOrEqualTo42) = list.partition { it < 42 }
println("lessThan42=$lessThan42")
println("greaterThanOrEqualTo42=$greaterThanOrEqualTo42")
Enter fullscreen mode Exit fullscreen mode

And the output is

lessThan42=[-5, 25]
greaterThanOrEqualTo42=[100, 42]
Enter fullscreen mode Exit fullscreen mode

However, for some reason, there is no partition() function for maps. So, let's right the wrong and create an extension function for it:

fun <K, V> Map<K, V>.partition(predicate: (Map.Entry<K, V>) -> Boolean): Pair<Map<K, V>, Map<K, V>> {
    val map1 = mutableMapOf<K, V>()
    val map2 = mutableMapOf<K, V>()
    entries.forEach {
        if (predicate(it)) {
            map1[it.key] = it.value
        } else {
            map2[it.key] = it.value
        }
    }
    return Pair(map1, map2)
}
Enter fullscreen mode Exit fullscreen mode

Here is the usage example:

val list = listOf(100, -5, 42, 25)
val map = list.associateWith { it.toString() }
val (keyLessThan42, keyGreaterThanOrEqualTo42) = map.partition { it.key < 42 }
println("keyLessThan42=$keyLessThan42")
println("keyGreaterThanOrEqualTo42=$keyGreaterThanOrEqualTo42")
Enter fullscreen mode Exit fullscreen mode

And the output is

keyLessThan42={-5=-5, 25=25}
keyGreaterThanOrEqualTo42={100=100, 42=42}
Enter fullscreen mode Exit fullscreen mode

Dream your code, code your dream.

Top comments (0)