Search code examples
indexingjuliavectorization

Julia - How to index vector w/ vector of vector of indices?


Suppose idxs is a vector of vectors defined as

3-element Vector{Vector{Int64}}:
 [1, 2]
 [1, 3]
 [1, 4]

And suppose A is a vector defined as

5-element Vector{Int64}:
 6
 7
 8
 9
 10

With Python and numpy, I could simply do (in pseudocode) A[idxs] and this would return a 2-D array containing the elements

[
  [6 7]
  [6 8]
  [6 9]
]

How do I perform the same indexing in Julia, where I am indexing multiple times into a 1-dimensional vector, and retrieving a matrix of entries?

Thanks


Solution

  • You can use an array comprehension with hcat. You can either use the resulting 2x3 Matrix or transpose the result.

    julia> hcat([A[i] for i in id]...)'
    3×2 adjoint(::Matrix{Int64}) with eltype Int64:
     6  7
     6  8
     6  9
    

    Data

    julia> show(id)
    [[1, 2], [1, 3], [1, 4]]
    
    julia> show(A)
    [6, 7, 8, 9, 10]