# Selecting entries of an array based on selection criteria provided by other array

**URL:** https://discourse.julialang.org/t/selecting-entries-of-an-array-based-on-selection-criteria-provided-by-other-array/118632
**Category:** New to Julia
**Tags:** performance, arrays, allocations, comprehension
**Created:** [August 26, 2024, 4:28pm UTC](https://discourse.julialang.org/t/selecting-entries-of-an-array-based-on-selection-criteria-provided-by-other-array/118632 "2024-08-26T16:28:25Z")
**Posts on this page:** 11
**Page:** 1

<div class="post-metadata">

### Author: ![oscarvdvelde](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/oscarvdvelde/32/202157_2.png) [@oscarvdvelde](https://discourse.julialang.org/u/oscarvdvelde)
#### Post date: [August 26, 2024, 4:28pm UTC](https://discourse.julialang.org/t/selecting-entries-of-an-array-based-on-selection-criteria-provided-by-other-array/118632/1 "2024-08-26T16:28:25Z")

</div>

Consider this example, which may run in a loop where `selection` is updated.

```julia
a = [1.0,1.2,0.9, 3.4, 4.5,4.5, 9.7, -1.5,-2.1]
b = [1,1,1,2,3,3,4,5,5] # e.g. a grouping of a
selection = [1,3,5] # these groups

# What I may naively want to do, and then e.g. median over it
julia> a[b.==selection]  
ERROR: DimensionMismatch: arrays could not be broadcast to a common size; got a dimension with lengths 9 and 3

# My probably naive solution:
vcat([a[b.==s] for s in selection]...)

```

It must be a recipe for an excess of allocations. Is there a better way?

---

<div class="post-metadata">

### Author: ![JonasWickman](https://avatars.discourse-cdn.com/v4/letter/j/9de0a6/32.png) [@JonasWickman](https://discourse.julialang.org/u/JonasWickman)
#### Post date: [August 26, 2024, 5:30pm UTC](https://discourse.julialang.org/t/selecting-entries-of-an-array-based-on-selection-criteria-provided-by-other-array/118632/2 "2024-08-26T17:30:38Z")

</div>

The most compact way to do this is probably

```julia
a[b .∈ (selection,)] # or a[b .∈ Ref(selection)]

```

but you have to use `∈` (type `\in[tab]` in REPL to get this symbol) in order to broadcast. The `(,)`or `Ref` means broadcast will treat `selection` as a scalar for broadcast.

This is not ideal for allocation/speed, in part because it creates an intermediate bit array. The fastest/least allocating solution I could come up with is not particularly pretty:

```julia
function myfun(a, b, selection)
    
    l = 0
    for bi ∈ b 
       if bi ∈ selection; l += 1; end
    end
    out = zeros(eltype(a), l)
    
    k = 1
    for i ∈ eachindex(a, b)
        if b[i] ∈ selection
            out[k] = a[i]
            k += 1
        end
    end
    
    return out
end

```

---

<div class="post-metadata">

### Author: ![stevengj](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/stevengj/32/71_2.png) [@stevengj](https://discourse.julialang.org/u/stevengj)
#### Post date: [August 26, 2024, 5:55pm UTC](https://discourse.julialang.org/t/selecting-entries-of-an-array-based-on-selection-criteria-provided-by-other-array/118632/3 "2024-08-26T17:55:11Z")

</div>

You could also do:

```julia
[a[i] for i in eachindex(a,b) if b[i] in selection]

```

(If efficiency matters, and you can rely on `b` and `selection` being sorted as in your example, then there are faster algorithms where you don’t have to repeat the `in` search for every element.)

---

<div class="post-metadata">

### Author: ![JonasWickman](https://avatars.discourse-cdn.com/v4/letter/j/9de0a6/32.png) [@JonasWickman](https://discourse.julialang.org/u/JonasWickman)
#### Post date: [August 26, 2024, 6:07pm UTC](https://discourse.julialang.org/t/selecting-entries-of-an-array-based-on-selection-criteria-provided-by-other-array/118632/4 "2024-08-26T18:07:00Z")

</div>

> [@stevengj](#):
>
> You could also do:
> 
> `[a[i] for i in eachindex(a,b) if b[i] in selection]`

I hope this isn’t derailing the thread too much, but I tried this and surprisingly it’s about half as fast as `a[b .∈ (selection,)]`on my machine for this specific example. My best guess is that it’s because the length of the array is not known in advance without first computing the intermediate bit array, but if anyone knows for sure I’d love to find out.

---

<div class="post-metadata">

### Author: ![oscarvdvelde](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/oscarvdvelde/32/202157_2.png) [@oscarvdvelde](https://discourse.julialang.org/u/oscarvdvelde)
#### Post date: [August 26, 2024, 9:07pm UTC](https://discourse.julialang.org/t/selecting-entries-of-an-array-based-on-selection-criteria-provided-by-other-array/118632/5 "2024-08-26T21:07:49Z")

</div>

Thanks, very useful. I will try it in my code. I like the one-liner, which can be memorized easily, but also the looping function could see good use, as I will probably be throwing this into functions regularly.

On my laptop, when placed in a function, the above code performs like:

```julia
vcat([a[b.==s] for s in selection]...)
281.848 ns (11 allocations: 720 bytes)
3.334 ms (62 allocations: 1.71 MiB) on actual data (see below)

```

```julia
a[b .∈ Ref(selection)]
77.137 ns (3 allocations: 208 bytes)
4.472 ms (5 allocations: 149.33 KiB) on actual data

```

```julia
loop function
60.285 ns (1 allocation: 112 bytes)
10.826 ms (1 allocation: 4.25 KiB) on actual data

```

```julia
[a[i] for i in eachindex(a,b) if b[i] in selection]
135.664 ns (3 allocations: 208 bytes)
4.537 ms (7 allocations: 11.12 KiB) on actual data

```

I also tried this shorter version of Jonas’ loop function:

```julia
function test5(a, b, selection)
   out = zeros(eltype(a), length(b))
   k = 1
   for i ∈ eachindex(a, b)
       if b[i] ∈ selection
           out[k] = a[i]
           k += 1
       end
   end
   filter!(x->x!=0.0,out)
   return out
end

56.548 ns (1 allocation: 128 bytes)
5.489 ms (3 allocations: 4.40 MiB) on actual data 
(allocates unnecessarily for entire vector, which is why Jonas' first loop is needed, but it is faster)

```

On the other hand, my vectors may not be sorted at all (sorted by time, while vector `a` could be latitude, not lined up with the order of `b` but the same length). I still have to figure out `eachindex(a,b)` and its uses.

**Update:** I have now listed also the performance when applied to one iteration of my actual data, where `a` and `b` are of length 1153405, and selection of length 12.

---

<div class="post-metadata">

### Author: ![oscarvdvelde](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/oscarvdvelde/32/202157_2.png) [@oscarvdvelde](https://discourse.julialang.org/u/oscarvdvelde)
#### Post date: [August 26, 2024, 11:35pm UTC](https://discourse.julialang.org/t/selecting-entries-of-an-array-based-on-selection-criteria-provided-by-other-array/118632/6 "2024-08-26T23:35:27Z")

</div>

See my updated numbers above! For large data arrays the relative results look quite different.

The two single line functions perform similarly, the Jonas loop takes twice as long, but allocates close to nothing, while my original list comprehension is surprisingly fast, allocating a lot, but still less than the single loop (the Jonas loop function I tried to simplify).

It seems easy to make loops, in hope of speed and allocation improvements, that instead allocate way more than needed, here 4 MB vs 4 kB! Good to consider including a pre-allocation loop from now on.

---

<div class="post-metadata">

### Author: ![GunnarFarneback](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gunnarfarneback/32/1827_2.png) [@GunnarFarneback](https://discourse.julialang.org/u/GunnarFarneback)
#### Post date: [August 27, 2024, 7:16am UTC](https://discourse.julialang.org/t/selecting-entries-of-an-array-based-on-selection-criteria-provided-by-other-array/118632/7 "2024-08-27T07:16:59Z")

</div>

> [@JonasWickman](#):
>
> The most compact way to do this is probably
> 
> ```julia
> a[b .∈ (selection,)] # or a[b .∈ Ref(selection)]
> 
> ```
> 
> but you have to use `∈` (type `\in[tab]` in REPL to get this symbol) in order to broadcast.

Those who are unicode-adverse can broadcast with

```julia
a[in.(b, (selection,))]

```

> [@oscarvdvelde](#):
>
> ```julia
> filter!(x->x!=0.0,out)
> 
> ```

This doesn’t work correctly if you have zeros in `a` (which may or may not be a problem for you) but since you know that you only want to discard data at the end, you can do it more efficiently with

```julia
resize!(out, k - 1)

```

> [@oscarvdvelde](#):
>
> For large data arrays the relative results look quite different.

The density of the selection definitely matters for the relative efficiency of these approaches.  
Edit: As does the size of `selection`. If it is big you certainly don’t want to check inclusion more than once per `b` element.

---

<div class="post-metadata">

### Author: ![GunnarFarneback](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gunnarfarneback/32/1827_2.png) [@GunnarFarneback](https://discourse.julialang.org/u/GunnarFarneback)
#### Post date: [August 27, 2024, 1:24pm UTC](https://discourse.julialang.org/t/selecting-entries-of-an-array-based-on-selection-criteria-provided-by-other-array/118632/8 "2024-08-27T13:24:23Z")

</div>

Here’s another way to spell the broadcasting:

```julia
a[in(selection).(b)]

```

If `selection` is large and the set of all classes is a range from 1, it may be interesting to try some variation of

```julia
binary_selection = in(selection).(1:maximum(b))
a[binary_selection[b]]

```

---

<div class="post-metadata">

### Author: ![nilshg](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/nilshg/32/2283_2.png) [@nilshg](https://discourse.julialang.org/u/nilshg)
#### Post date: [August 28, 2024, 4:27pm UTC](https://discourse.julialang.org/t/selecting-entries-of-an-array-based-on-selection-criteria-provided-by-other-array/118632/9 "2024-08-28T16:27:06Z")

</div>

Or `in(b).(selection)`

---

<div class="post-metadata">

### Author: ![GunnarFarneback](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gunnarfarneback/32/1827_2.png) [@GunnarFarneback](https://discourse.julialang.org/u/GunnarFarneback)
#### Post date: [August 28, 2024, 7:38pm UTC](https://discourse.julialang.org/t/selecting-entries-of-an-array-based-on-selection-criteria-provided-by-other-array/118632/10 "2024-08-28T19:38:11Z")

</div>

That gives a quite different result though, of little use for this problem.

---

<div class="post-metadata">

### Author: ![nilshg](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/nilshg/32/2283_2.png) [@nilshg](https://discourse.julialang.org/u/nilshg)
#### Post date: [August 28, 2024, 8:35pm UTC](https://discourse.julialang.org/t/selecting-entries-of-an-array-based-on-selection-criteria-provided-by-other-array/118632/11 "2024-08-28T20:35:36Z")

</div>

Ah sorry ignore, I just wanted to point to the Fix1 version of in, hadn’t seen that you already noted that in your later post (quite apart from the fact that I confused the argument order…)
