How do you access multi-dimension array by N array of index element-wise?

user_5269577 0
Viewed 137

Suppose we have

A = [1 2; 3 4]

In numpy, the following syntax will produce

A[[1,2],[1,2]] = [1,4]

But, in julia, the following produce a permutation which output

A[[1,2],[1,2]] = [1 2; 3 4]

Is there a concise way to achieve the same thing as numpy without using for loops?

julia
2 Answers

To get what you want I would use CartesianIndex like this:

julia> A[CartesianIndex.([(1,1), (2,2)])]
2-element Vector{Int64}:
 1
 4

or

julia> A[[CartesianIndex(1,1), CartesianIndex(2,2)]]
2-element Vector{Int64}:
 1
 4
user_12695670

Like Bogumil said, you probably want to use CartesianIndex. But if you want to get your result from supplying the vectors of indices for each dimensions, as in your Python [1,2],[1,2] example, you need to zip these indices first:

julia> A[CartesianIndex.(zip([1,2], [1,2]))]
2-element Vector{Int64}:
 1
 4

How does this work? zip traverses both vectors of indices at the same time (like a zipper) and returns an iterator over the tuples of indices:

julia> zip([1,2],[1,2]) # is a lazy iterator
zip([1, 2], [1, 2])

julia> collect(zip([1,2],[1,2])) # collect to show all the tuples
2-element Vector{Tuple{Int64, Int64}}:
 (1, 1)
 (2, 2)

and then CartesianIndex turns them into cartesian indices, which can then be used to get the corresponding values in A:

julia> CartesianIndex.(zip([1,2],[1,2]))
2-element Vector{CartesianIndex{2}}:
 CartesianIndex(1, 1)
 CartesianIndex(2, 2)
user_56021490
Related