How do you perform conditional assignment in arrays in Julia?

Question

In Octave, I can do

octave:1> A = [1 2; 3 4]
A =

   1   2
   3   4

octave:2> A(A>1) -= 1
A =

   1   1
   2   3

but in Julia, the equivalent syntax does not work.

julia> A = [1 2; 3 4]
2x2 Array{Int64,2}:
 1  2
 3  4

julia> A[A>1] -= 1
ERROR: `isless` has no method matching isless(::Int64, ::Array{Int64,2})
 in > at operators.jl:33

How do you conditionally assign values to certain array or matrix elements in Julia?

Answer 1

Your problem isn't with the assignment, per se, it's that A > 1 itself doesn't work. You can use the elementwise A .> 1 instead:

julia> A = [1 2; 3 4];

julia> A .> 1
2×2 BitArray{2}:
 false  true
  true  true

julia> A[A .> 1] .-= 1000;

julia> A
2×2 Array{Int64,2}:
    1  -998
 -997  -996

Update:

Note that in modern Julia (>= 0.7), we need to use . to say that we want to broadcast the action (here, subtracting by the scalar 1000) to match the size of the filtered target on the left. (At the time this question was originally asked, we needed the dot in A .> 1 but not in .-=.)

Answer 2

To make it work in Julia 1.0 one need to change = to .=. In other words:

julia> a = [1 2 3 4]

julia> a[a .> 1] .= 1

julia> a
1×4 Array{Int64,2}:
 1  1  1  1

Otherwise you will get something like

ERROR: MethodError: no method matching setindex_shape_check(::Int64, ::Int64)

Answer 3

In Julia v1.0 you can use the replace! function instead of logical indexing, with considerable speedups:

julia> B = rand(0:20, 8, 2);

julia> @btime (A[A .> 10] .= 10) setup=(A=copy($B))
  595.784 ns (11 allocations: 4.61 KiB)

julia> @btime replace!(x -> x>10 ? 10 : x, A) setup=(A=copy($B))
  13.530 ns ns (0 allocations: 0 bytes)

For larger matrices, the difference hovers around 10x speedup.

The reason for the speedup is that the logical indexing solution relies on creating an intermediate array, while replace! avoids this.

A slightly terser way of writing it is

replace!(x -> min(x, 10), A)

There doesn't seem to be any speedup using min, though.

And here's another solution that is almost as fast:

A .= min.(A, 10)

and that also avoids allocations.