Killing a For loop in Julia array comprehension

问题

I have the following line of code in Julia:

X=[(i,i^2) for i in 1:100 if i^2%5==0]

Basically, it returns a list of tuples (i,i^2) from i=1 to 100 if the remainder of i^2 and 5 is zero. What I want to do is, in the array comprehension, break out of the for loop if i^2 becomes larger than 1000. However, if I implement

X=[(i,i^2) for i in 1:100 if i^2%5==0 else break end]

I get the error: syntax: expected "]".

Is there any way to easily break out of this for loop inside the array? I've tried looking online, but nothing came up.

回答1:

I don't think so. You could always just

tmp(i) = (j = i^2; j > 1000 ? false : j%5==0)
X=[(i,i^2) for i in 1:100 if tmp(i)]

回答2:

It's a "fake" for-loop, so you can't break it. Take a look at the lowered code below:

julia> foo() = [(i,i^2) for i in 1:100 if i^2%5==0]
foo (generic function with 1 method)

julia> @code_lowered foo()
LambdaInfo template for foo() at REPL[0]:1
:(begin 
        nothing
        #1 = $(Expr(:new, :(Main.##1#3)))
        SSAValue(0) = #1
        #2 = $(Expr(:new, :(Main.##2#4)))
        SSAValue(1) = #2
        SSAValue(2) = (Main.colon)(1,100)
        SSAValue(3) = (Base.Filter)(SSAValue(1),SSAValue(2))
        SSAValue(4) = (Base.Generator)(SSAValue(0),SSAValue(3))
        return (Base.collect)(SSAValue(4))
    end)

The output shows that array comprehension is implemented via Base.Generator which takes an iterator as input. It only supports the [if cond(x)::Bool] "guard" for now, so there is no way to use break here.

For your specific case, a workaround is to use isqrt:

julia> X=[(i,i^2) for i in 1:isqrt(1000) if i^2%5==0]
6-element Array{Tuple{Int64,Int64},1}:
 (5,25)  
 (10,100)
 (15,225)
 (20,400)
 (25,625)
 (30,900)

回答3:

Using a for loop is considered idiomatic in Julia and could be more readable in this instance. Also, it could be faster.

Specifically:

julia> using BenchmarkTools

julia> tmp(i) = (j = i^2; j > 1000 ? false : j%5==0)
julia> X1 = [(i,i^2) for i in 1:100 if tmp(i)];

julia> @btime [(i,i^2) for i in 1:100 if tmp(i)];
  471.883 ns (7 allocations: 528 bytes)

julia> X2 = [(i,i^2) for i in 1:isqrt(1000) if i^2%5==0];

julia> @btime [(i,i^2) for i in 1:isqrt(1000) if i^2%5==0];
  281.435 ns (7 allocations: 528 bytes)

julia> function goodsquares()
           res = Vector{Tuple{Int,Int}}()
           for i=1:100
               if i^2%5==0 && i^2<=1000
                   push!(res,(i,i^2))
               elseif i^2>1000
                   break
               end
           end
           return res
       end
julia> X3 = goodsquares();

julia> @btime goodsquares();
  129.123 ns (3 allocations: 304 bytes)

So, another 2x improvement is nothing to disregard and the long function gives plenty of room for illuminating comments.

来源：https://stackoverflow.com/questions/44097910/killing-a-for-loop-in-julia-array-comprehension

标签

arrays

julia

list-comprehension

break