Search code examples
kotlinlist-comprehensiondictionary-comprehension

What is the equivalent of Python list, set, and map comprehensions in Kotlin?

In Python, there are list comprehensions and similar constructs for maps and sets. In Kotlin there is nothing at all in any of the documentation with a similar name.

What are the equivalents of these comprehensions? For example, those found in Python 3 Patterns, Recipes and Idioms. Which includes comprehensions for:

  • list
  • set
  • dictionary

Note: this question is intentionally written and answered by the author (Self-Answered Questions), so that the idiomatic answers to commonly asked Kotlin topics are present in SO.

Solution

  • Taking examples from Python 3 Patterns, Recipes and Idioms we can convert each one to Kotlin using a simple pattern. The Python version of a list comprehension has 3 parts:

    1. output expression
    2. input list/sequence and variable
    3. optional predicate

    These directly correlate to Kotlin functional extensions to collection classes. The input sequence, followed by the optional predicate in a filter lambda, followed by the output expression in a map lambda. So for this Python example:

    # === PYTHON

a_list = [1, 2, 3, 4, 5, 6]

#                    output | var | input   | filter/predicate 
even_ints_squared = [ e*e  for e in a_list  if e % 2 == 0 ]

print(even_ints_squared)
# output: [ 4, 16, 36 ]

    Becomes

    // === KOTLIN

var aList = listOf(1, 2, 3, 4, 5, 6)

//                    input      |   filter      |       output              
val evenIntsSquared = aList.filter { it % 2 == 0 }.map { it * it }

println(evenIntsSquared)
// output: [ 4, 16, 36 ]

    Notice that the variable is not needed in the Kotlin version since the implied it variable is used within each lambda. In Python you can turn these into a lazy generator by using the () instead of square brackets:

    # === PYTHON

even_ints_squared = ( e**2 for e in a_list if e % 2 == 0 )

    And in Kotlin it is more obviously converted to a lazy sequence by changing the input via a function call asSequence():

    // === KOTLIN

val evenIntsSquared = aList.asSequence().filter { it % 2 == 0 }.map { it * it }

    Nested comprehensions in Kotlin are created by just nesting one within the other's map lambda. For example, take this sample from PythonCourse.eu in Python changed slightly to use both a set and a list comprehension:

    # === PYTHON

noprimes = {j for i in range(2, 8) for j in range(i*2, 100, i)}
primes = [x for x in range(2, 100) if x not in noprimes]
print(primes)
# output: [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97]

    Becomes:

    // === KOTLIN

val nonprimes = (2..7).flatMap { (it*2..99).step(it).toList() }.toSet()
val primes = (2..99).filterNot { it in nonprimes }
print(primes)    
// output: [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79, 83, 89, 97]

    Notice that the nested comprehension produces a list of lists which is converted to a flat list using flatMap() and then converted to a set using toSet(). Also, Kotlin ranges are inclusive, whereas a Python range is exclusive so you will see the numbers are slightly different in the ranges.

    You can also use a sequence generator with co-routines in Kotlin to yield the values without needing the call to flatMap() or flatten():

    // === KOTLIN

val nonprimes = sequence {
    (2..7).forEach { (it*2..99).step(it).forEach { value -> yield(value) } }
}.toSet()
val primes = (2..99).filterNot { it in nonprimes }

    Another example from the referenced Python page is generating a matrix:

    # === PYTHON

matrix = [ [ 1 if item_idx == row_idx else 0 for item_idx in range(0, 3) ] for row_idx in range(0, 3) ]
print(matrix)
# [[1, 0, 0], 
#  [0, 1, 0], 
#  [0, 0, 1]]

    And in Kotlin:

    // === KOTLIN

val matrix = (0..2).map { row -> (0..2).map { col -> if (col == row) 1 else 0 }}
println(matrix)
// [[1, 0, 0], 
//  [0, 1, 0], 
//  [0, 0, 1]]

    Or in Kotlin instead of lists, you could also generate arrays:

    // === KOTLIN

val matrix2 = Array(3) { row -> 
                          IntArray(3) { col -> if (col == row) 1 else 0 } 
                       }

    Another of the examples for set comprehensions is to generate a unique set of properly cased names:

    # === PYTHON

names = [ 'Bob', 'JOHN', 'alice', 'bob', 'ALICE', 'J', 'Bob' ]

fixedNames = { name[0].upper() + name[1:].lower() for name in names if len(name) > 1 }

print(fixedNames)
# output: {'Bob', 'Alice', 'John'}

    Is translated to Kotlin:

    // === KOTLIN

val names = listOf( "Bob", "JOHN", "alice", "bob", "ALICE", "J", "Bob" )

val fixedNames = names.filter { it.length > 1 }
       .map { it.take(1).toUpperCase() + it.drop(1).toLowerCase() }
       .toSet()

println(fixedNames)
// output: [Bob, John, Alice]

    And the example for map comprehension is a bit odd, but can also be implemented in Kotlin. The original:

    # === PYTHON

mcase = {'a':10, 'b': 34, 'A': 7, 'Z':3}

mcase_frequency = { k.lower() : mcase.get(k.lower(), 0) + mcase.get(k.upper(), 0) for k in mcase.keys() }

print(mcase_frequency)
# output: {'a': 17, 'z': 3, 'b': 34}

    And the converted, which is written to be a bit more "wordy" here to make it clearer what is happening:

    // === KOTLIN

val mcase = mapOf("a" to 10, "b" to 34, "A" to 7, "Z" to 3)

val mcaseFrequency = mcase.map { (key, _) ->
    val newKey = key.toLowerCase()
    val newValue = mcase.getOrDefault(key.toLowerCase(), 0) +
                   mcase.getOrDefault(key.toUpperCase(), 0)
    newKey to newValue
}.toMap()

print(mcaseFrequency)
// output: {a=17, b=34, z=3}

    Further reading: