I would like to take an array x and change all numbers greater than 5 to 5. What is the standard way to do this in one line?
Below is some code that does this in several lines. This question on logical indexing is related but appears to concern selection rather than assignment. Thanks
x = [1 2 6 7]
for i in 1:length(x)
if x[i] >= 5
x[i] = 5
end
end
Desired output: x = [1 2 5 5]
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.
Collected from the Internet
Please contact [email protected] to delete if infringement.
Comments