Skip to content

Instantly share code, notes, and snippets.

@heinrisch
Last active May 9, 2016 10:39
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save heinrisch/d6ceac8fb373a880496c4a7a2e1cf8ad to your computer and use it in GitHub Desktop.
Save heinrisch/d6ceac8fb373a880496c4a7a2e1cf8ad to your computer and use it in GitHub Desktop.
Easy lookup creation from lists
import java.util.*
// Already in Kotlin std lib: list.mapTo(mutableSetOf(), keyGenerator)
fun <K, V> List<V>.createLookupSet(keyGenerator: (V) -> K) : HashSet<K> {
return fold(HashSet<K>(), {set, value -> set.add(keyGenerator(value)); set})
}
// Already in Kotlin std lib: list.associateBy(keyGenerator)
fun <K, V> List<V>.createLookupMap(keyGenerator: (V) -> K) : HashMap<K, V> {
return fold(HashMap<K, V>(), {map, value -> map.put(keyGenerator(value), value); map})
}
// Already in Kotlin std lib: list.groupBy(keyGenerator)
fun <K, V> List<V>.createLookupMapList(keyGenerator: (V) -> K) : HashMap<K, MutableList<V>> {
return DefaultListHashMap<K, V>().createLookup(this, keyGenerator)
}
// Instead of this, use getOrPut when wanting to insert into HashMap for same functionality
// Example: hashMap.getOrPut(key) { mutableListOf() }
class DefaultListHashMap<K, V>() : HashMap<K, MutableList<V>>() {
fun put(key: K, value: V) {
check(key)
get(key)!!.add(value)
}
fun check(key: K) {
if(!containsKey(key)){
put(key, mutableListOf())
}
}
override fun get(key: K): MutableList<V> {
check(key)
return super.get(key)!!
}
fun createLookup(values : List<V>, keyGenerator: (V) -> K) : DefaultListHashMap<K, V> {
values.forEach { put(keyGenerator(it), it) }
return this
}
}
@heinrisch
Copy link
Author

Sample usage:

val storeLookup = listOfMyStores.createLookupMapList { it.merchant.id }
myMerchants.forEach { merchant -> 
    val merchantStores = storeLookup.get(merchant.id)
    // do more things...
}

@ilya-g
Copy link

ilya-g commented May 8, 2016

These functions are already in Kotlin Standard Library. They are called mapTo, associateBy, groupBy respectively.

val merchants = listOfMyStores.mapTo(mutableSetOf()) { it.merchant.id } 
// also one could use listOfMyStores.map { it.merchant.id }.toSet()

val storeById = listOfMyStores.associateBy { it.id }
val storeByMerchantId = listOfMyStores.groupBy { it.merchant.id }

@heinrisch
Copy link
Author

heinrisch commented May 8, 2016

@ilya-g Perfect, thanks! I'll update the gist

Is there such thing as a DefaultDict also?
Edit: Found getOrPut by looking at the source of groupBy, that should work for a defaultDict!

@hastebrot
Copy link

hastebrot commented May 9, 2016

SSCCE (Short, Self Contained, Correct Example) for this:

data class Store(val id: Int, val merchant: Merchant)
data class Merchant(val id: String)

fun main(args: Array<String>) {
    val merchant1 = Merchant("foo")
    val merchant2 = Merchant("bar")
    val stores = listOf(
        Store(1, merchant1),
        Store(2, merchant2),
        Store(3, merchant2)
    )

    val merchants = stores.mapTo(mutableSetOf()) { it.merchant.id }
//    val merchants = stores.map { it.merchant.id }.toSet()
    val storeById = stores.associateBy { it.id }
    val storeByMerchantId = stores.groupBy { it.merchant.id }

    println(merchants)
    println(storeById)
    println(storeByMerchantId)
}

Update: Example for MutableMap<K, V>.getOrPut() which is like Python's dict.setdefault() or defaultdict, but additionally inlined and using a lazy default value.

val map = mutableMapOf<String, MutableList<Int>>()
map.getOrPut("foo") { mutableListOf() }.add(1)
map.getOrPut("bar") { mutableListOf() }.add(2)
map.getOrPut("bar") { mutableListOf() }.add(3)
println(map) // {foo=[1], bar=[2, 3]}

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment