Commit 275b1c4f authored by Éric Thiébaut's avatar Éric Thiébaut
Browse files

Provide select and select!

parent 05aa6a76
......@@ -231,6 +231,83 @@ end
pupil_shape(A::AbstractMatrix{Bool}) = A
#------------------------------------------------------------------------------
# SELECTION OF ENTRIES
"""
Tao.select(A, sel)
yields the values of array `A` for which the predicate `sel(A[i])` is true.
Argument `sel` may also be a boolean array indicating which elements of `A` to
select.
This method is equivalent to `filter(sel,A)` except for the order of the
arguments and for the possibility to use a boolean mask.
"""
function select(A::AbstractArray{T,N},
sel::AbstractArray{Bool,N}) where {T,N}
axes(A) == axes(sel) || dimension_mismatch(
"arguments `A` and `sel` must have the same indices")
# FIXME: count(sel) is as fast as the following code
n = 0
@inbounds for i in eachindex(sel)
n += (sel[i] ? 1 : 0)
end
return select!(Array{T}(undef, n), A, sel)
end
function select(A::AbstractArray, sel)
n = 0
@inbounds for i in eachindex(A)
n += (sel(A[i]) ? 1 : 0)
end
return select!(Array{eltype(A)}(undef, n), A, sel)
end
"""
Tao.select!(dst, A, sel) -> dst
overwrites the contents of `dst` with the values of array `A` for which the
predicate `sel(A[i])` is true. Argument `sel` may also be a boolean array
indicating which elements of `A` to select. It is asserted that the length of
the destination `dst` is exactly the number of selected elements.
"""
function select!(dst::AbstractVector{<:Any},
A::AbstractArray{<:Any,N},
sel::AbstractArray{Bool,N}) where {N}
axes(A) == axes(sel) || dimension_mismatch(
"arguments `A` and `sel` must have the same indices")
J = Base.axes1(dst)
j = first(J) - 1
n = last(J)
@inbounds for i in eachindex(A, sel)
if sel[i]
j += 1
j n || error("too many selected values")
dst[j] = A[i]
end
end
j == n || error("too few selected values")
return dst
end
function select!(dst::AbstractVector, A::AbstractArray, sel)
J = Base.axes1(dst)
j = first(J) - 1
n = last(J)
@inbounds for i in eachindex(A)
if sel(A[i])
j += 1
j n || error("too many selected values")
dst[j] = A[i]
end
end
j == n || error("too few selected values")
return dst
end
#------------------------------------------------------------------------------
# VECTORIZED OPERATIONS
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment