CODEWITHSUNDEEP
kotlinkey-valuemultimap
dcompiled
dcompiled

Reputation: 4832

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.

Upvotes: 4

Views: 2284

Answers (3)

WildM
WildM

Reputation: 135

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.

Upvotes: 0

Sweeper
Sweeper

Reputation: 271575

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
    }

Upvotes: 5

lukas.j
lukas.j

Reputation: 7173

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]}

Upvotes: 0

Related Questions