kotlinkey-valuemultimap

Kotlin transform List<Pair<K, Collection<V>>> into Multimap


I'm looking for an idiomatic way of converting a list of pairs where Pair.first is a key and Pair.second is a list of values. This procedural approach works but I was hoping to find a more idiomatic way that doesn't require creating the mutable lists directly.

val pairs: Pair<String, List<Int>>

val res = mutableMapOf<String, List<Int>>()
pairs.forEach {
    res.getOrPut(it.first, ::mutableListOf).addAll(it.second)
}

This code can get wrapped in an extension function like follows but it doesn't seem very generic:

fun <K, V> List<Pair<K, Collection<V>>>.toMultimap(): Map<K, List<V>> {
    var res = mutableMapOf<K, MutableList<V>>()
    forEach {
        res.getOrPut(it.first, ::mutableListOf).addAll(it.second)
    }
    return res
}

Using pairs.toMap doesn't work because it overwrites map keys with a last in wins approach. groupBy works comes close, it creates keys to values in a list of lists structure.

val pairs2 = listOf(
    Pair("a", listOf(1, 2, 3)),
    Pair("b", listOf(6, 7)),
    Pair("a", listOf(4, 5)),
    Pair("b", listOf(8, 9)),
)

val res = pairs2.groupBy({ it.first }, { it.second })
println(res)

{a=[[1, 2, 3], [4, 5]], b=[[6, 7], [8, 9]]}

It is possible to then flatten the map but the downside here is that its pretty inefficient as this creates double the required hashmaps and lists per key (one for groupby and another for flatten). If there

val res = pairs2.groupBy({ it.first }, { it.second }).mapValues { it.value.flatten() }
println(res)

{a=[1, 2, 3, 4, 5], b=[6, 7, 8, 9]}

Looking to see if there are any better approaches to accomplishing this transform.


Solution

  • Rather than groupBy, use groupingBy, which produces a Grouping. This is an intermediate object on which you can do all kinds of fold/reduce operations. In your case:

    fun <K, V> List<Pair<K, Collection<V>>>.toMultimap() =
        groupingBy { it.first }
            .fold(emptyList<V>()) { acc, (_, new) -> acc + new }
    

    If you don't like the fact that + creates too many new lists, you can do something like this:

    groupingBy { it.first }
        .fold({ _, _ -> mutableListOf<V>() }) { _, acc, (_, new) ->
            acc.addAll(new)
            acc
        }