# Fast method to select matrix elements based on a vector of positions

There is an operation for which I have long wanted to find a better solution.

Let:

• a be a matrix of dimensions $m\times n$
• v be an integer vector of length $n$ with elements drawn from $[1, m]$

For every element $x$ at position $p$ in v I wish to select the element at row $x$, column $p$ in a.

Example:

SeedRandom[0]

a = Array[Range[7] 10^# &, 3, 0]

v = RandomInteger[{1,3}, 7]


$\left( \begin{array}{ccccccc} 1 & 2 & 3 & 4 & 5 & 6 & 7 \\ 10 & 20 & 30 & 40 & 50 & 60 & 70 \\ 100 & 200 & 300 & 400 & 500 & 600 & 700 \\ \end{array} \right)$

{3, 3, 2, 1, 1, 3, 1}


Desired output:

{100, 200, 30, 4, 5, 600, 7}


### Details:

• Although a compiled function is likely to be the fastest approach for packed arrays I want something more general, allowing arrays of mixed types, and ideally optimized for arrays in which each row is a packed array (list) of a different type, e.g.

$\left( \begin{array}{ccccccc} 1 & 2 & 3 & 4 & 5 & 6 & 7 \\ 0.1 & 0.2 & 0.3 & 0.4 & 0.5 & 0.6 & 0.7 \\ \text{a} & \text{b} & \text{c} & \text{d} & \text{e} & \text{f} & \text{g} \\ \end{array} \right)$

• I am still interested in seeing the fastest possible compiled function as it may serve as the basis for a general solution as well.

• I seek a solution that works well for any shape of array a, from $n\gg m$ to square to $m\gg n$, though if compromise is necessary I would optimize for $n > m$.

-

This seems fast(er):

Extract[a, Transpose[{v, Range@Length@v}]]


Mr.Wizard's clean method

Diagonal @ a[[v]]


has a surprising property for those of us who think that packed arrays rank just below the wheel in the list of inventions for the sake of efficiency. For unpacked arrays a, it uses virtually no extra memory.

Example

Initialization. Start with a fresh kernel. After quitting is finished, initialize a and v.

Quit[]

SeedRandom[0]
a = RandomReal[1, {50, 10000}];
v = RandomInteger[{1, 50}, 10000];


Example 1. In the case a is a packed array, a[[v]] is computed and fills memory.

MaxMemoryUsed[]
Diagonal@a[[v]]; // AbsoluteTiming
MaxMemoryUsed[]
(*
32251456
{0.420822, Null}
831474792
*)


Example 2. Below unpacked is not packed. Evidently filling memory with unpacked[[v]] is deferred and Diagonal is computed without expanding unpacked[[v]] further.

unpacked = a; unpacked[[1, 1]] = "x";

MaxMemoryUsed[]
Diagonal@unpacked[[v]]; // AbsoluteTiming
MaxMemoryUsed[]
(*
55228232
{0.001294, Null}
55566776
*)


Example 3. (Edit) There is a way to take some advantage of this to improve the performance of Diagonal on packed array by partial unpacking and using Tr. See Mr.Wizard's answer, under Update.

For comparison purposes with examples 1 and 2, here are my timings for the Extract method:

First@AbsoluteTiming@ Do[Extract[a, Transpose[{v, Range@Length@v}]], {100}] /100
(*
0.00009192
*)

First@AbsoluteTiming@ Do[Extract[unpacked, Transpose[{v, Range@Length@v}]], {100}] /100
(*
0.00057431
*)

-
This is just the sort of thing I was looking for. Funny how easy it seems now but so do many great ideas. Thanks –  Mr.Wizard Jul 26 at 17:37

Although I am self-answering, as stated, I am not satisfied with these approaches.
Nevertheless they may be useful and they can serve as a benchmark for any new solutions.

This is cleanest method I know, though sadly it is a true memory hog, and not fast either:

Diagonal @ a[[v]]

{100, 200, 30, 4, 5, 600, 7}


More practically but still slow:

MapThread[Part, {a\[Transpose], v}]


Even slower:

MapIndexed[a[[#, First @ #2]] &, v]


Auto-compilation in Table makes it much faster for fully packed arrays but elsewhere it is slower than MapThread, including mixed arrays where every row is packed.

Table[a[[ v[[i]], i ]], {i, Length@v}]


### Observation:

Notably only the first approach uses v directly without iterating over it, and it is impractical due to the huge expression created. I wonder if I have overlooked a way to perform a vectorized extraction without that memory explosion.

## Update

Michael astutely observed:

If unpacked is not a packed array, Diagonal@unpacked[[v]] is pretty fast and uses almost no extra memory.

Looking at this again I realize there are two problems:

1. When extracting rows from a fully packed array the expressions are not memory-shared

2. Diagonal fully unpacks

The first problem can be avoided by unpacking the outermost level of the array, and the second by using a function that does not unpack: Tr.

Combined:

(List @@ a)[[v]] ~Tr~ List

{100, 200, 30, 4, 5, 600, 7}


As an alternative to List @@ a we can use the seemingly inert replacement h_[x__] :> h[x] to unpack the outer level. This leads to a rather fun definition:

h_[x__] ~f~ v_ := h[x][[v]] ~Tr~ h


This revised method actually performs quite well, though it is still half as fast as Extract on unpacked data and more than an order of magnitude behind on a packed array.

-
If unpacked is not a packed array, Diagonal@unpacked[[v]] is pretty fast and uses almost no extra memory. –  Michael E2 Jul 26 at 17:25
@MichaelE2 Great observation! I'm so glad I posted this question. Why don't you put that in your answer too, so it doesn't get lost? –  Mr.Wizard Jul 26 at 17:40
@MichaelE2 Please see Update –  Mr.Wizard Jul 27 at 12:03
I thought about exploring unpacking just the first level, but I didn't have time. I doubt I would have come up with Tr, and I'm sure I wouldn't have come up with such a cute definition. It would be worth another upvote, if I could. :) –  Michael E2 Jul 27 at 13:10

This could be fast.

  a[[v[[#]], #]] & /@ Range[Length@v];

-
I haven't yet found a case where this is particularly good. In a single test on fully packed data it is not as fast as Table, and on unpacked data it is slower than Table and significantly slower than MapThread. Have you found a case where this is faster than the options in my own answer? –  Mr.Wizard Jul 26 at 17:45
I checked it with this: SeedRandom[0] a = RandomReal[100, {13000, 13000}]; v = RandomInteger[{1, 13000}, 13000];. I got zero Timing. –  Algohi Jul 26 at 18:19
For me Table still performs better with that data but this is not far off and it's another option. +1 –  Mr.Wizard Jul 27 at 12:07

You could do it with a sparse array:

s = SparseArray[

But sadly, Flatten will take a long time for large a. (If you could keep around Flatten[a] for many "queries", it might be competitive, though).
That is unlike anything I tried. I do like the idea of using Dot. With some data you will find that Join @@ a is faster than Flatten[a, 1]. Thanks for a novel idea. +1 –  Mr.Wizard Jul 27 at 12:10
Your method inspired this idea: Flatten[a, 1][[Subtract[v, 1] # + Range[#] & @ Length @ v]]. It's pretty good but still not as fast as Extract. –  Mr.Wizard Jul 27 at 12:49