Slice a map given a list of keys in Kotlin

Question

Given a map, and a list of keys

val abc = mapOf(1 to "a", 2 to "b", 3 to "c")
val keys = listOf(1, 2)

How do I a get a map containing only the key-value pairs specified by keys? Something like

val ab = abc.slice(keys)
// equivalent to mapOf(1 to "a", 2 to "b)

I'm looking for something a bit more elegant than

val ab = listOf(1, 2).map { it to abc[it] }.toMap()

For example, in Elixir:

abc = %{1 => "a", 2 => "b", 3 => "c"}
ab = Map.take(abc, [1, 2])
# equivalent to ab = %{1 => "a", 2 => "b"}

Dextrin

Answer 1

You can use filterKeys:

val ab = abc.filterKeys { it in keys }

And since it is Kotlin, you could even define your own extension function to achieve exactly what you imagined:

fun <T> Map<T, *>.slice(keys: Iterable<T>) = filterKeys { it in keys }

val ab = abc.slice(keys)

Rakes

it in keys is not a good idea for large lists, this makes the entire operation to be O(mXn) – Earley

FilterKeys is super-slow, no idea, why. – Nowise

Answer 2

Solutions given in above answers does solve the problem but I think a small change is warranted.

Problem is that for every key in the map they check if the list contains that key, which is O(n) operation, this is ok for small lists but once you reach a certain size it becomes very slow. I suggest that you convert the list of keys to a set which reduces the contains operation to O(1) in average case. (Hence reducing carbon footprint :) ).

Following is the solution with above change incorporated.

val mapAbc = mapOf(1 to "a", 2 to "b", 3 to "c")
val keySet = listOf(1, 2).toSet()
val filteredMap = mapAbc.filterKeys { it in keySet }

Earley

It's no more efficient or fewer steps than listOf(1, 2).map { it to abc[it] }.toMap() from the question, though 😎. It does feel slightly more idomatic, but in most cases I wouldn't bother with the set conversion, it's something to remember every time. Feels like this is missing from the standard library. – Dextrin

my bad, I overlooked that part. but the point is that contains operation on a list is expensive and if you need to perform multiple of them, then you need to reconsider your data structures. things always start small and before you know the list that should have stored only 100 elements end up storing 10000 items, making your runtime less than optimal – Earley

Actually, .filterKeys(x in keys) is brutally (!) slow for large lists, tested it myself. – Nowise

Answer 3

abc.filterKeys { it in listOf(1, 2) }

Chiliarch

Answer 4

fun <K, V> Map<K, V>.slice(keys: Collection<K>): Map<K,V> {
    val resultMap = mutableMapOf<K, V>()
    for (key in keys) {
        get(key)?.let{resultMap[key] = it}
    }
    return resultMap
}

Test:

@Test
fun `should slice map by keys`() {
    // GIVEN

    val inputMap = (1..10_000).associate {
        UUID.randomUUID() to "value $it"
    }

    val nonExistingKey = UUID.randomUUID()
    val filterKeys = inputMap.keys.take(999).shuffled() + nonExistingKey

    // WHEN

    val startSlice = System.currentTimeMillis()
    val filteredMapSlice = inputMap.slice(filterKeys)
    println("Duration of slice: ${System.currentTimeMillis() - startSlice} ms")

    val startFilterKeys = System.currentTimeMillis()
    val filteredMapFilterKeys = inputMap.filterKeys { it in filterKeys }
    println("Duration of filterKeys: ${System.currentTimeMillis() - startFilterKeys} ms")
    assertThat(filteredMapFilterKeys).isEqualTo(filteredMapSlice)

    // THEN
    
assertThat(filteredMapSlice).hasSize(filterKeys.size - 1) // non-existing key should have been filtered
    assertThat(filteredMapSlice.keys).doesNotContain(nonExistingKey)
    assertThat(filteredMapSlice.keys).allMatch { it in inputMap.keys }
    filteredMapSlice.forEach{ (key, value) ->
        assertThat(value).isEqualTo(inputMap[key])
    }
}

This is magnitudes (!) faster than .filterKeys()

Duration of slice: 3 ms
Duration of filterKeys: 1479 ms

Nowise

How does it compare to the suggestion in the question? – Dextrin

Answer 5

A minor variation of @spyro's efficient answer:

fun <K, V> Map<K, V>.slice(keys: Collection<K>): Map<K,V> =
    buildMap {
        keys.forEach { key ->
            this@slice[key]?.let { value ->
                put(key, value)
            }
        }
    }

Bejarano