indexing - How to change array elements that are greater than 5 to 5, in one line?

Question

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]

question from:https://stackoverflow.com/questions/65878950/how-to-change-array-elements-that-are-greater-than-5-to-5-in-one-line

Answer 1

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.