diff --git a/base/array.jl b/base/array.jl index bd9d3b87335416..b3e0c3625d61a3 100644 --- a/base/array.jl +++ b/base/array.jl @@ -2336,20 +2336,46 @@ Int64[] function findall(A) collect(first(p) for p in pairs(A) if last(p)) end + # Allocating result upfront is faster (possible only when collection can be iterated twice) -function findall(A::AbstractArray{Bool}) - n = count(A) +function findall(f::Function, A::AbstractArray{Bool}) + # Compute f for true and false only once + ft, ff = f(true), f(false) + (ft | ff) || return Vector{eltype(keys(A))}() + (ft & ff) && return vec(Array(keys(A))) + n = let + c = count(A) + ft ? c : length(A) - c + end I = Vector{eltype(keys(A))}(undef, n) + _findall(ff, I, A) +end + +function _findall(invert::Bool, I::Vector, A::AbstractArray{Bool}) cnt = 1 - for (i,a) in pairs(A) - if a - I[cnt] = i - cnt += 1 - end + len = length(I) + for (k, v) in pairs(A) + cnt > len && break + I[cnt] = k + cnt += v ⊻ invert end I end +function _findall(invert::Bool, I::Vector, A::AbstractVector{Bool}) + i = firstindex(A) + cnt = 1 + len = length(I) + @inbounds while cnt ≤ len + I[cnt] = i + cnt += A[i] ⊻ invert + i = nextind(A, i) + end + I +end + +findall(A::AbstractArray{Bool}) = findall(identity, A) + findall(x::Bool) = x ? [1] : Vector{Int}() findall(testf::Function, x::Number) = testf(x) ? [1] : Vector{Int}() findall(p::Fix2{typeof(in)}, x::Number) = x in p.x ? [1] : Vector{Int}() diff --git a/test/arrayops.jl b/test/arrayops.jl index f2905aa0b582a2..071338a71c0355 100644 --- a/test/arrayops.jl +++ b/test/arrayops.jl @@ -547,9 +547,17 @@ end @testset "findall, findfirst, findnext, findlast, findprev" begin a = [0,1,2,3,0,1,2,3] + m = [false false; true false] @test findall(!iszero, a) == [2,3,4,6,7,8] @test findall(a.==2) == [3,7] @test findall(isodd,a) == [2,4,6,8] + @test findall(Bool[]) == Int[] + @test findall([false, false]) == Int[] + @test findall(m) == [k for (k,v) in pairs(m) if v] + @test findall(!, [false, true, true]) == [1] + @test findall(i -> true, [false, true, false]) == [1, 2, 3] + @test findall(i -> false, rand(2, 2)) == Int[] + @test findall(!, m) == [k for (k,v) in pairs(m) if !v] @test findfirst(!iszero, a) == 2 @test findfirst(a.==0) == 1 @test findfirst(a.==5) == nothing