CODEWITHSUNDEEP
arraysjulia
Ben S.
Ben S.

Reputation: 3545

Given an index of choices for each column, construct a 1D array from a 2D array

I have a 2D array such as:

julia> m = [1 2 3 4 5
            6 7 8 9 10
            11 12 13 14 15]
3×5 Array{Int64,2}:
  1   2   3   4   5
  6   7   8   9  10
 11  12  13  14  15

I want to pick one value from each column and construct a 1D array.

So for instance, if my choices are 

julia> choices = [1, 2, 3, 2, 1]
5-element Array{Int64,1}:
 1
 2
 3
 2
 1

Then the desired output is [1, 7, 13, 9, 5]. What's the best way to do that? In my particular application, I am randomly generating these values, e.g.

choices = rand(1:size(m)[1], size(m)[2])

Thank you!

Upvotes: 2

Views: 43

Answers (1)

Bogumił Kamiński
Bogumił Kamiński

Reputation: 69899

This is probably the simplest approach: [m[c, i] for (i, c) in enumerate(choices)]

EDIT: If best means fastest for you such a function should be approximately 2x faster than the comprehension for large m:

function selector(m, choices)
    v = similar(m, size(m, 2))
    for i in eachindex(choices)
        @inbounds v[i] = m[choices[i], i]
    end
    v
end

Upvotes: 3

Related Questions