Faster array indexing with a logical array

Hi everyone,

Pretty basic question…

I was experimenting with indexing arrays with logical arrays. For instance, say I have an array shaped (10000,2) and I need to extract all the indices where the sum across the columns is less than a threshold. What would be the best Julian way to do it? I tried list comprehension but that is slower than the method below

MWE:

using BenchmarkTools 

# what I could think of ?
function getindicesVectorized(x::Array)
    return x[getindex.(findall(sum(x, dims=2) .<= 1), 1), :];
end

function getindicesListComp(x::Array)
    ind = findall(sum(x, dims=2) .<= 1)
    temp = [x[ind[i][1], :] for i in 1:length(ind)]
    return hcat(temp...)';
end

points = rand(10000,2)
julia> @btime getindicesVectorized($points)
  32.201 μs (15 allocations: 279.70 KiB)

It seems to be fast and readable (see equivalent numpy version), but I was wondering if there would be a faster way to do this ?

What I would do in python naively

import numpy as np
from timeit import timeit
points = np.random.random((10000,2))
In [8]: %timeit points[np.sum(points, axis=1) <= 1, :]
212 µs ± 2.73 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

Thanks in advance !!

12 posts - 6 participants

Read full topic