# Iterating over several arrays

Hi!
Say I have 2 arrays:

``````a=rand(3);
b=rand(4);
``````

and I want to iterate over both of them

``````for i in 1:7
do stuff on ab[i]
end
``````

what’s a good way of doing that? The obvious way is `ab = vcat(a,b)`, but that is creating a new array (my example is small, my real case is large).
Any clever way of doing it without having to create a new array and without a bunch of if statements to switch from `a` to `b` inside the loop (and, obviously, without 2 for loops)?
Thanks!

I do not understand the question. Could you elaborate on the `do stuff` part or provide a MWE ?

``````for thing in Iterators.flatten([a, b])
... do sth with thing
end
``````
1 Like

Any way to do this with no allocations? Also, is there the option to output the index as well? I can always build the index myself, of course.
Here’s an example:

``````a=[1,2,3]
b=[4,5]
function tmpfun(a,b)
acc=0
for element in Iterators.flatten([a,b])
acc+=element
end
return acc
end
``````

Thanks a lot!

``````function tmpfun2(a,b)
acc=0
for v in (a,b)
for e in v
acc+=e
end
end
acc
end
``````
2 Likes

I think there is something in the Iterators module, `chain` or something.

You can avoid the allocation using `Iterators.flatten((a, b))` instead of `Iterators.flatten([a, b])`.

3 Likes

My bad, use a tuple literal

``````julia> @benchmark sum(thing for thing in Iterators.flatten((A,B))) setup = begin A = rand(10); B = rand(20) end
BenchmarkTools.Trial:
memory estimate:  0 bytes
allocs estimate:  0
--------------
minimum time:     14.113 ns (0.00% GC)
median time:      14.315 ns (0.00% GC)
mean time:        14.414 ns (0.00% GC)
maximum time:     38.939 ns (0.00% GC)
--------------
samples:          10000
evals/sample:     999

julia> @benchmark sum(thing for thing in Iterators.flatten([A,B])) setup = begin A = rand(10); B = rand(20) end
BenchmarkTools.Trial:
memory estimate:  96 bytes
allocs estimate:  1
--------------
minimum time:     41.994 ns (0.00% GC)
median time:      44.209 ns (0.00% GC)
mean time:        47.920 ns (4.97% GC)
maximum time:     1.501 μs (96.88% GC)
--------------
samples:          10000
evals/sample:     993
``````

Edit2:

Also, is there the option to output the index as well?

Just the running number? You can wrap an `enumerate` around it. Or do you mean the indices into `a` and `b`? That would be a tiny bit more code, but you probably wouldn’t want to concatenate them in that case. It’s easily solved by sudete. I’d say we have a draw here sudete ``````julia> for (i,thing) in enumerate(Iterators.flatten((a,b)))
# do something useful here
print(i," ",thing," ")
end
1 0.11576933628769859 2 0.5420240722242764 3 0.2224035318833757 [...]
``````
3 Likes

yep barely If you want a single index sequence from 1 to the total number, you can use `for (i, x) in enumerate(Iterators.flatten((a, b)))`.

If you want the original indices from the arrays, you can do

``````function f(a,b)
for (i, x) in Iterators.flatten((pairs(a), pairs(b)))
println("\$i: \$x")
end
end

julia> f([1,2,3], [40,50])
1: 1
2: 2
3: 3
1: 40
2: 50
``````

This also works with non-standard arrays with indices starting at 0 or whatever, and it doesn’t allocate:

``````function f(a,b)
acc_i, acc_x = 0, 0
for (i, x) in Iterators.flatten((pairs(a), pairs(b)))
acc_i += i
acc_x += x
end
return (acc_i, acc_x)
end

julia> @btime f(\$[1,2,3], \$[40, 50])
15.846 ns (0 allocations: 0 bytes)
``````

Edit: This time @FPGro was faster for `enumerate` 3 Likes

Thanks a lot, folks. This is what I ended up using:

``````function tmpfun(a,b)
acc=0;
acc2=0
for (i,ii) in enumerate(Iterators.flatten((a,b)))
acc+=i;
acc2+=ii
end
return acc,acc2
end
``````

which does exactly what I needed (global index, no allocations).
Thanks again!

1 Like

And in case anyone is curious, I compared `tmpfun` to

``````function tmpfun2(a,b)
acc=0;
acc2=0
for i in 1:length(a)
acc+=i;
acc2+=a[i];
end
for i in 1:length(b)
acc+=i+length(a);
acc2+=b[i];
end
return acc,acc2
end
``````

and:

``````a=rand(10000);
b=rand(5000);

julia> @btime tmpfun(\$a,\$b)
18.399 μs (3 allocations: 64 bytes)
(112507500, 7565.072862093092)

julia> @btime tmpfun2(\$a,\$b)
17.999 μs (3 allocations: 64 bytes)
(112507500, 7565.072862093092)
``````

So it seems like there is no performance penalty!
The allocations happen when I return `acc,acc2`. If I return `acc+acc2` there are no allocations. So it’s allocating the array that gets returned, I guess (no clue as to why it says 3 allocations, but that’s not really an issue).
Thanks again!

These allocations are due to a type instability: `acc2` starts as an integer but is later assigned a float. You can see this by running `@code_warntype tmpfun2(a,b)`.

To fix it you can initialize `acc2=0.0`, or to work with any element type: `acc2=zero(eltype(a))`.

If `a` and `b` can have different element types, you would need `zero(promote_type(eltype(a), eltype(b)))` but you would get lots of allocations from `Iterators.flatten` anyway…

3 Likes

Ah, I see. Thanks for the tips!

If you just add up the indices from 1 to n, you don’t even need to keep track of them explicitly. There’s a simple formula:

``````function bar(iters)
s = sum(length, iters)
return (s*(s+1)÷2, sum(Iterators.flatten(iters)))
end
``````

Which is faster, more general and free of allocations ^^

``````julia> @benchmark tmpfun(A,B) setup = begin A = rand(10000); B = rand(5000) end
BenchmarkTools.Trial:
memory estimate:  64 bytes
allocs estimate:  3
--------------
minimum time:     13.699 μs (0.00% GC)
median time:      13.900 μs (0.00% GC)
mean time:        14.047 μs (0.00% GC)
maximum time:     65.700 μs (0.00% GC)
--------------
samples:          10000
evals/sample:     1

julia> @benchmark bar((A,B)) setup = begin A = rand(10000); B = rand(5000) end
BenchmarkTools.Trial:
memory estimate:  0 bytes
allocs estimate:  0
--------------
minimum time:     10.199 μs (0.00% GC)
median time:      10.300 μs (0.00% GC)
mean time:        10.396 μs (0.00% GC)
maximum time:     44.699 μs (0.00% GC)
--------------
samples:          10000
evals/sample:     1
``````

Yep, I’m familiar with Gauss’ formula. The accumulator was just a demo. My real code is way more complicated, but requires the indices as well (to access other vectors that are as long as a+b).
Thanks!

1 Like