How to use if-else in a list comprehension

user_10892426 0
Viewed 76

I wanted to compare execution time for two list comprehensions one which uses if-else keywords and another that uses ternary operator.

arr = [5, 8, 12, 17, 24, 42];

function squares_and_cubes(array)
    return [ x%2 ==0 ?  x^2 : x^3 for x in array]
end

function squares_and_cubes_if_else(array)
    [ x^2 if (x%2==0) else x^3 end for x in array]
end
@time squares_and_cubes(arr)
@time squares_and_cubes_if_else(arr)

However, I get the following error message ERROR: syntax: invalid comprehension syntax when trying to execute [ x^2 if (x%2==0) else x^3 end for x in arr]. As a Julia newbie, I can't figure out what is the issue with this syntax - I am using v 1.7 .

julia if-statement list-comprehension
3 Answers

It seems you're trying to use Python-like if-else syntax here. In Julia it is:

julia> arr = [1, 2, 3, 4];

julia> [if (x%2==0) x^2 else x^3 end for x in arr]
4-element Vector{Int64}:
  1
  4
 27
 16

There should be no difference in time of execution since both get parsed into the pretty much the same abstract code underneath, and will generate the exact same machine code. (The ternary operator is more readable and idiomatic for this purpose though.)

But if you did want to indulge your curiosity, make sure to run the functions once before timing them so as to avoid measuring the compilation time.

sundar_81270

There are many interesting forms in Julia that does this than in Python, some of which don't need array comprehension. You can use ifelse, map or even function broadcasting; all having similar performance.

@btime [ifelse(x%2 == 0, x^2, x^3) for x in $arr] # 39.376 ns (1 allocation: 112 bytes)
@btime [x%2 == 0 ? x^2 : x^3 for x in $arr]       # 39.879 ns (1 allocation: 112 bytes)
@btime map(x -> x%2 == 0 ? x^2 : x^3, $arr)       # 39.678 ns (1 allocation: 112 bytes)
@btime (x -> x%2 == 0 ? x^2 : x^3).($arr)         # 36.858 ns (1 allocation: 112 bytes)
user_39431700

Another option is to use the ternary operator which results in a shorter code:

julia> arr = [1, 2, 3, 4];

julia> [x % 2 == 0 ? x^2 : x^3 for x in arr]
4-element Vector{Int64}:
  1
  4
 27
 16
user_99577100
Related