Using an array for in array indexing

问题

Is there a nicer way to do this?

a = rand(10,10)
a[[CartesianIndex(i,i) for i in 1:10]] .= something

I.e. using setindex! without a CartesianIndex or something like

a[i,i for i in 1:10] .= something

Or in general using setindex! with an array for syntax?

回答1:

I like the CartesianIndex solution, but one other way would be to create a boolean mask:

julia> a = rand(5, 5);

julia> a[[i == j for i ∈ 1:5, j ∈ 1:5]] .= 0;

julia> a
5×5 Array{Float64,2}:
 0.0       0.376169  0.0248078  0.957535   0.522565
 0.116614  0.0       0.743684   0.0500792  0.704349
 0.728995  0.643491  0.0        0.367225   0.348022
 0.338079  0.451293  0.383365   0.0        0.111781
 0.602485  0.249625  0.963363   0.0850144  0.0

回答2:

I just wanted to sum up all the answers and comments looking at the performance of each proposed solution.

using LinearAlgebra, BenchmarkTools
A = rand(1000, 1000)
B  = [i == j for i ∈ 1:1000, j ∈ 1:1000]
B2 = BitArray(B)
CI = CartesianIndex.(1:1000, 1:1000)
step = size(A,1) +1
    
@btime A[[i == j for i ∈ 1:1000, j ∈ 1:1000]]
2.622 ms (3 allocations: 984.64 KiB)

@btime A[B]
1.806 ms (1 allocation: 7.94 KiB)

@btime A[(1:1000) .== (1:1000)']
509.597 μs (5 allocations: 134.38 KiB)

@btime A[B2]
35.067 μs (1 allocation: 7.94 KiB)

@btime A[[CartesianIndex(i,i) for i in 1:1000]]
6.051 μs (2 allocations: 23.69 KiB)

@btime A[CartesianIndex.(1:1000, 1:1000)]
5.769 μs (2 allocations: 23.69 KiB)

@btime A[CI]
4.093 μs (1 allocation: 7.94 KiB)

@btime A[1:size(A,1)+1:end]
2.123 μs (4 allocations: 8.00 KiB)

@btime A[1:step:end]
2.073 μs (3 allocations: 7.98 KiB)

@btime A[diagind(A)]
2.036 μs (2 allocations: 7.97 KiB)

Using linear indices is the fastest solution, where the inbuilt diagind performs slightly better. Accessing out of diagonal elements is probably easier with CartesianIndices and still performs quite good.

来源：https://stackoverflow.com/questions/65041097/using-an-array-for-in-array-indexing

标签

indexing

julia