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

Just for exercise the closest to python will be:

infix fun <I, O> ((I) -> O).`in`(range: Iterable<I>): List<O> = range.map(this).toList()
infix fun <I> Iterable<I>.`if`(cond: (I) -> Boolean): List<I> = this.filter(cond)

fun main() {
    { it: Int -> it + 1 } `in` 1..2 `if` {it > 0}
}

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:

Tags:

List Comprehension

Dictionary Comprehension

Kotlin

Recent Posts