add replace & replace! for collections

base/algorithms.jl

@@ -0,0 +1,119 @@
+module Algorithms
+
+export replace!, replace
+
+import Base: replace
+
+
+"""
+    replace!(pred, [f::Function], A, [new], [count])
+
+Replace all occurrences `x` in collection `A` for which `pred(x)` is true
+by `new` or `f(x)` (exactly one among `f` and `new` must be specified).
+If `count` is specified, then replace at most `count` occurrences.
+This is the in-place version of [`replace`](@ref).
+
+# Examples
+```jldoctest
+julia> a = [1, 2, 3, 1];
+
+julia> replace!(isodd, a, 0, 2); a
+4-element Array{Int64,1}:
+ 0
+ 2
+ 0
+ 1
+
+julia> replace!(x->x.first=>3, Dict(1=>2, 3=>4), 1) do (k, v)
+           v < 3
+       end
+
+Dict{Int64,Int64} with 2 entries:
+  3 => 4
+  1 => 3
+```
+
+!!! note
+When `A` is an `Associative` or `AbstractSet` collection, if
+there are collisions among old and newly created keys, the result
+can be unexpected:
+
+```jldoctest
+julia> replace!(x->true, x->2x, Set([3, 6]))
+Set([12])
+```
+"""
+function replace!(pred, new::Function, A::AbstractArray, n::Integer=-1)
+    n == 0 && return A
+    count = 0
+    @inbounds for i in eachindex(A)
+        if pred(A[i])
+            A[i] = new(A[i])
+            count += 1
+            count == n && break
+        end
+    end
+    A
+end
+
+askey(k, ::Associative) = k.first
+askey(k, ::AbstractSet) = k
+
+function replace!(pred, new::Function, A::Union{Associative,AbstractSet}, n::Integer=-1)
+    n == 0 && return A
+    del = eltype(A)[]
+    count = 0
+    for x in A
+        if pred(x)
+            push!(del, x)
+            count += 1
+            count == n && break
+        end
+    end
+    for k in del
+        pop!(A, askey(k, A))
+        push!(A, new(k))
+    end
+    A
+end
+
+const ReplaceCollection = Union{AbstractArray,Associative,AbstractSet}
+
+replace!(pred, A::ReplaceCollection, new, n::Integer=-1) = replace!(pred, y->new, A, n)
+
+"""
+    replace(pred, [f::Function], A, [new], [count])
+
+Return a copy of collection `A` where all occurrences `x` for which
+`pred(x)` is true are replaced by `new` or `f(x)` (exactly one among
+`f` and `new` must be specified).
+If `count` is given, then replace at most `count` occurrences.
+See the in-place version [`replace!`](@ref) for examples.
+"""
+replace(pred, new::Function, A::ReplaceCollection, n::Integer=-1) = replace!(pred, new, copy(A), n)
+replace(pred, A::ReplaceCollection, new, n::Integer=-1) = replace!(pred, copy(A), new, n)
+
+"""
+    replace!(A, old, new, [count])
+
+Replace all occurrences of `old` in collection `A` by `new`.
+If `count` is given, then replace at most `count` occurrences.
+
+# Examples
+```jldoctest
+julia> replace!(Set([1, 2, 3]), 1, 0)
+Set([0, 2, 3])
+```
+"""
+replace!(A::ReplaceCollection, old, new, n::Integer=-1) = replace!(x->x==old, A, new, n)
+
+"""
+    replace(A, old, new, [count])
+
+Return a copy of collection `A` where all occurrences of `old` are
+replaced by `new`.
+If `count` is given, then replace at most `count` occurrences.
+"""
+replace(A::ReplaceCollection, old, new, n::Integer=-1) = replace!(copy(A), old, new, n)
+
+end

base/exports.jl

@@ -1351,6 +1351,9 @@ export
     nzrange,
     nnz,
 
+# Algorithms module re-exports
+    replace!,
+
 # Distributed module re-exports
     @spawn,
     @spawnat,

base/sysimg.jl

@@ -406,6 +406,10 @@ import .Dates: Date, DateTime, DateFormat, @dateformat_str, now
 include("sparse/sparse.jl")
 importall .SparseArrays
 
+# algorithms
+include("algorithms.jl")
+importall .Algorithms
+
 include("asyncmap.jl")
 
 include("distributed/Distributed.jl")

test/algorithms.jl

@@ -0,0 +1,28 @@
+
+@testset "replace! & replace" begin
+    a = [1, 2, 3, 1]
+    @test replace(iseven, x->2x, a) == [1, 4, 3, 1]
+    @test replace!(iseven, x->2x, a) === a
+    @test a == [1, 4, 3, 1]
+    @test replace(a, 1, 0) == [0, 4, 3, 0]
+    @test replace(a, 1, 0, 1) == [0, 4, 3, 1] # 1 replacement only
+    @test replace!(a, 1, 2) == [2, 4, 3, 2]
+
+    d = Dict(1=>2, 3=>4)
+    @test replace(x->x.first > 2, d, 0=>0) == Dict(1=>2, 0=>0)
+    @test replace!(x->x.first > 2, x->(x.first=>2*x.second), d) ==
+        Dict(1=>2, 3=>8)
+    @test replace(d, 3=>8, 0=>0) == Dict(1=>2, 0=>0)
+    @test replace!(d, 3=>8, 2=>2) === d
+    @test d == Dict(1=>2, 2=>2)
+    @test replace(x->x.second == 2, d, 0=>0, 1) in [Dict(1=>2, 0=>0),
+                                                    Dict(2=>2, 0=>0)]
+
+    s = Set([1, 2, 3])
+    @test replace(x->x>1, x->2x, s) == Set([1, 4, 6])
+    @test replace(x->x>1, x->2x, s, 1) in [Set([1, 4, 3]), Set([1, 2, 6])]
+    @test replace(s, 1, 4) == Set([2, 3, 4])
+    @test replace!(s, 1, 2) == Set([2, 3])
+
+    @test !(0 in replace([1, 2, 3], 1, 0, 0)) # count=0 --> no replacements
+end