Collection processing in Kotlin: Associating elements

This is a chapter from the book Functional Kotlin. You can find it on LeanPub or Amazon. It is also available as a course.

// associate simplified implementation from Kotlin stdlib
inline fun <T, K, V> Iterable<T>.associate(
    transform: (T) -> Pair<K, V>
): Map<K, V> {
    val destination = LinkedHashMap<K, V>(initialCapacity())
    for (element in this) {
        destination += transform(element)
    }
    return destination
}

// associateBy simplified implementation from Kotlin stdlib
inline fun <T, K> Iterable<T>.associateBy(
    keySelector: (T) -> K
): Map<K, T> {
    val destination = LinkedHashMap<K, V>(initialCapacity())
    for (element in this) {
        destination.put(keySelector(element), element)
    }
    return destination
}

// associateWith simplified implementation from Kotlin stdlib
public inline fun <K, V> Iterable<K>.associateWith(
    valueSelector: (K) -> V
): Map<K, V> {
    val destination = LinkedHashMap<K, V>(initialCapacity())
    for (element in this) {
        destination.put(element, valueSelector(element))
    }
    return destination
}

To transform an iterable[^087_0] into a map, we use the associate method. In maps, elements are represented by both a key and a value, therefore the associate method needs to return a pair. If you want to use the elements of your list as the keys of your new map, a better alternative to associate is associateWith. On its lambda expression, you should specify what the value should be for each key. If you want to use elements of your list as values of your new map, a better alternative to associate is associateBy. On its lambda expression, specify what the key should be for each value.

fun main() {
    val names = listOf("Alex", "Ben", "Cal")
    println(names.associate { it.first() to it.drop(1) })
    // {A=lex, B=en, C=al}
    println(names.associateWith { it.length })
    // {Alex=4, Ben=3, Cal=3}
    println(names.associateBy { it.first() })
    // {A=Alex, B=Ben, C=Cal}
}

associateWith(op) works the same as associate { it to op(it) }. associateBy(op) works the same as associate { op(it) to it }.

Be careful because keys on maps need to be unique, and a new value with the same key replaces the previous one. If you want to keep instead of replace previous values, use the groupBy or groupingBy method instead of the associateBy method.

fun main() {
    val names = listOf("Alex", "Aaron", "Ada")
    println(names.associateBy { it.first() })
    // {A=Ada}
    println(names.groupBy { it.first() })
    // {A=[Alex, Aaron, Ada]}
}

When keys are unique, associateWith can be reversed using the keys property, and associateBy can be reversed using the values property.

fun main() {
    val names = listOf("Alex", "Ben", "Cal")
    val aW = names.associateWith { it.length }
    println(aW.keys.toList() == names) // true
    val aB = names.associateBy { it.first() }
    println(aB.values.toList() == names) // true
}

toList is required before comparison because keys returns a set, and values returns a custom collection, so both are represented with different collection types.

Finding an element in a list requires iterating over the elements one by one. Finding a value by a key is much more efficient thanks to the hash table that is used under the hood. That is why associateBy is used to optimize searching for elements[^087_1].

fun produceUserOffers(
    offers: List<Offer>,
    users: List<User>
): List<UserOffer> {
    //
    val usersById = users.associateBy { it.id }
    return offers
        .map { createUserOffer(it, usersById[it.buyerId]) }
}

distinct and distinctBy

// distinct implementation from Kotlin stdlib
fun <T> Iterable<T>.distinct(): List<T> {
    return this.toMutableSet().toList()
}

inline fun <T, K> Iterable<T>.distinctBy(
    selector: (T) -> K
): List<T> {
    val set = HashSet<K>()
    val list = ArrayList<T>()
    for (e in this) {
        val key = selector(e)
        if (set.add(key))
            list.add(e)
    }
    return list
}

So, we now know that we can use associate to transform a list to a map. Transforming it to a set is much easier: we can just use the toSet function. A set is much more similar to a list than a map, and the key difference is that sets do not allow duplicates[^087_2].

fun main() {
    val list: List<Int> = listOf(1, 2, 4, 2, 3, 1)
    val set: Set<Int> = list.toSet()
    println(set) // [1, 2, 4, 3]
}

If you want to keep operating on a list but at the same time eliminate duplicates, use the distinct method. Under the hood, it transforms a list into a set and then back to a list. So, it eliminates elements that are equal to each other.

fun main() {
    val numbers = listOf(1, 2, 4, 2, 3, 1)
    println(numbers) // [1, 2, 4, 2, 3, 1]
    println(numbers.distinct()) // [1, 2, 4, 3]

    val names = listOf("Marta", "Maciek", "Marta", "Daniel")
    println(names) // [Marta, Maciek, Marta, Daniel]
    println(names.distinct()) // [Marta, Maciek, Daniel]
}

We can also use distinctBy, which uses a selector and keeps only the elements with the distinct values returned by this selector. This way, it gives us full control over the criteria used to decide if two values are distinct.

fun main() {
    val names = listOf("Marta", "Maciek", "Marta", "Daniel")
    println(names) // [Marta, Maciek, Marta, Daniel]
    println(names.distinctBy { it[0] }) // [Marta, Daniel]
    println(names.distinctBy { it.length }) // [Marta, Maciek]
}

Be aware that distinct keeps the first element of the list, while associateBy keeps the last element.

fun main() {
    val names = listOf("Marta", "Maciek", "Daniel")
    println(names)
    // [Marta, Maciek, Daniel]
    println(names.distinctBy { it.length })
    // [Marta, Maciek]
    println(names.associateBy { it.length }.values)
    // [Marta, Daniel]
}

These functions are often used when we suspect that we accidentally have some kind of duplicates.

data class Person(val id: Int, val name: String) {
    override fun toString(): String = "$id: $name"
}

fun main() {
    val people = listOf(
        Person(0, "Alex"),
        Person(1, "Ben"),
        Person(1, "Carl"),
        Person(2, "Ben"),
        Person(0, "Alex"),
    )
    println(people.distinct())
    // [0: Alex, 1: Ben, 1: Carl, 2: Ben]
    println(people.distinctBy { it.id })
    // [0: Alex, 1: Ben, 2: Ben]
    println(people.distinctBy { it.name })
    // [0: Alex, 1: Ben, 1: Carl]
}

[^087_0]: I hope it is clear, that List and Set are iterables, because they implement Iterable interface.
[^087_1]: This optimization is better explained in Effective Kotlin, Item 55: Consider associating elements to a map.
[^087_2]: The second difference is that a set does not necessarily keep elements in order.