Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add keepat! for in-place logical filtering #42351

Merged
merged 1 commit into from
Sep 23, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 17 additions & 22 deletions base/abstractarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3067,29 +3067,9 @@ end

## keepat! ##

"""
keepat!(a::AbstractVector, inds)

Remove the items at all the indices which are not given by `inds`,
and return the modified `a`.
Items which are kept are shifted to fill the resulting gaps.

`inds` must be an iterator of sorted and unique integer indices.
See also [`deleteat!`](@ref).
# NOTE: since these use `@inbounds`, they are actually only intended for Vector and BitVector

!!! compat "Julia 1.7"
This function is available as of Julia 1.7.

# Examples
```jldoctest
julia> keepat!([6, 5, 4, 3, 2, 1], 1:2:5)
3-element Vector{Int64}:
6
4
2
```
"""
function keepat!(a::AbstractVector, inds)
function _keepat!(a::AbstractVector, inds)
local prev
i = firstindex(a)
for k in inds
Expand All @@ -3106,3 +3086,18 @@ function keepat!(a::AbstractVector, inds)
deleteat!(a, i:lastindex(a))
return a
end

function _keepat!(a::AbstractVector, m::AbstractVector{Bool})
length(m) == length(a) || throw(BoundsError(a, m))
j = firstindex(a)
for i in eachindex(a, m)
@inbounds begin
if m[i]
i == j || (a[j] = a[i])
j = nextind(a, j)
end
end
end
deleteat!(a, j:lastindex(a))
return a
end
48 changes: 48 additions & 0 deletions base/array.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2602,6 +2602,54 @@ function filter!(f, a::AbstractVector)
return a
end

"""
keepat!(a::Vector, inds)

Remove the items at all the indices which are not given by `inds`,
and return the modified `a`.
Items which are kept are shifted to fill the resulting gaps.

`inds` must be an iterator of sorted and unique integer indices.
See also [`deleteat!`](@ref).

!!! compat "Julia 1.7"
This function is available as of Julia 1.7.

# Examples
```jldoctest
julia> keepat!([6, 5, 4, 3, 2, 1], 1:2:5)
3-element Vector{Int64}:
6
4
2
```
"""
keepat!(a::Vector, inds) = _keepat!(a, inds)

"""
keepat!(a::Vector, m::AbstractVector{Bool})

The in-place version of logical indexing `a = a[m]`. That is, `keepat!(a, m)` on
vectors of equal length `a` and `m` will remove all elements from `a` for which
`m` at the corresponding index is `false`.

# Examples
```jldoctest
julia> a = [:a, :b, :c];

julia> keepat!(a, [true, false, true])
2-element Vector{Symbol}:
:a
:c

julia> a
2-element Vector{Symbol}:
:a
:c
```
"""
keepat!(a::Vector, m::AbstractVector{Bool}) = _keepat!(a, m)

# set-like operators for vectors
# These are moderately efficient, preserve order, and remove dupes.

Expand Down
3 changes: 3 additions & 0 deletions base/bitarray.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1063,6 +1063,9 @@ function deleteat!(B::BitVector, inds::AbstractVector{Bool})
return B
end

keepat!(B::BitVector, inds) = _keepat!(B, inds)
keepat!(B::BitVector, inds::AbstractVector{Bool}) = _keepat!(B, inds)

function splice!(B::BitVector, i::Integer)
# TODO: after deprecation remove the four lines below
# as v = B[i] is enough to do both bounds checking
Expand Down
20 changes: 20 additions & 0 deletions test/arrayops.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1457,6 +1457,26 @@ end
@test isempty(eoa)
end

@testset "logical keepat!" begin
# Vector
a = Vector(1:10)
keepat!(a, [falses(5); trues(5)])
@test a == 6:10
@test_throws BoundsError keepat!(a, trues(1))
@test_throws BoundsError keepat!(a, trues(11))

# BitVector
ba = rand(10) .> 0.5
@test isa(ba, BitArray)
keepat!(ba, ba)
@test all(ba)

# empty array
ea = []
keepat!(ea, Bool[])
@test isempty(ea)
end

@testset "deleteat!" begin
for idx in Any[1, 2, 5, 9, 10, 1:0, 2:1, 1:1, 2:2, 1:2, 2:4, 9:8, 10:9, 9:9, 10:10,
8:9, 9:10, 6:9, 7:10]
Expand Down