How to change array elements that are greater than 5 to 5, in one line?
You can broadcast min
as well:
x .= min.(x, 5)
Note that this is (slightly) more efficient than using x[x .> 5] .= 5
because it does not allocate the temporary array of Booleans, x .> 5
, and it can be automatically vectorized, with a single pass over the memory (as per Oscar's comment below):
julia> using BenchmarkTools
julia> x = [1 2 6 7] ; @btime $x .= min.($x, 5) ; # fast, no allocations
19.144 ns (0 allocations: 0 bytes)
julia> x = [1 2 6 7] ; @btime $x[$x .> 5] .= 5 ; # slower, allocates
148.678 ns (5 allocations: 304 bytes)
The broadcast operator .
works with any function, including relational operators, and it also works with assignment. Hence an intuitive one-liner is:
x[x .> 5] .= 5
This part x .> 5
broadcasts > 5
over x
, resulting in a vector of booleans indicating elements greater than 5. This part .= 5
broadcasts the assignment of 5
across all elements indicated by x[x .> 5]
.
However, inspired by the significant speed-up in Benoit's very cool answer below (please do check it out) I decided to also add an optimized variant with a speed test. The above approach, while very intuitive looking, is not optimal because it allocates a temporary array of booleans for the indices. A (more) optimal approach that avoids temporary allocation, and as a bonus will work for any predicate (conditional) function is:
function f_cond!(x::Vector{Int}, f::Function, val::Int)
@inbounds for n in eachindex(x)
f(x[n]) && (x[n] = val)
end
return x
end
So using this function we would write f_cond!(x, a->a>5, 5)
which assigns 5
to any element for which the conditional (anonymous) function a->a>5
evaluates to true
. Obviously this solution is not a neat one-liner, but check out the following speed tests:
julia> using BenchmarkTools
julia> x1 = rand(1:10, 100);
julia> x2 = copy(x1);
julia> @btime $x1[$x1 .> 5] .= 5;
327.862 ns (8 allocations: 336 bytes)
julia> @btime f_cond!($x2, a->a>5, 5);
15.067 ns (0 allocations: 0 bytes)
This is just ludicrously faster. Also, you can just replace Int
with T<:Any
. Given the speed-up, one might wonder if there is a function in Base
that already does this. A one-liner is:
map!(a->a>5 ? 5 : a, x, x)
and while this significantly speeds up over the first approach, it falls well short of the second.
Incidentally, I felt certain this must be a duplicate to another StackOverflow question, but 5 minutes searching didn't reveal anything.