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

Question

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.

Answer 1

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
}

Answer 2

val pairs: List<Pair<String, List<Int>>> = listOf(
Pair("a", listOf(1, 2, 3)),
Pair("b", listOf(6, 7)),
Pair("a", listOf(4, 5)),
Pair("b", listOf(8, 9)),
)
val result = pairs
.map { it.first }.distinct()
.associateWith { first -> pairs.filter { it.first == first }.flatMap { it.second } }
println(result) // {a=[1, 2, 3, 4, 5], b=[6, 7, 8, 9]}

Answer 3

Here's one way you could do it.
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.flatMap { pair ->
pair.second.map { value ->
pair.first to value
}
}.groupBy({ it.first }, { it.second })
println(res) // {a=[1, 2, 3, 4, 5], b=[6, 7, 8, 9]}
Essentially, you flatten the list in each pair, preserving the original key, and then you convert the whole thing into your Map.
If you're worried about the lists taking up too much space, you could instead work with sequences, though maybe that's overkill.
val res = pairs2.asSequence().flatMap { pair ->
pair.second.asSequence().map { value ->
pair.first to value
}
}.groupBy({ it.first }, { it.second })
The core idea here is that flatMap is excellent at allowing you to add/remove elements from your original list on the fly.