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

Question

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$.

Answer 1

This seems fast(er):

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

---

Addendum

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
*)

Answer 2

Edit: please see Update below.

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.

Answer 3

You could do it with a sparse array:

s = SparseArray[
   MapThread[({#1, #2} -> 1) &,
    {
     Range[Length[v]],
     (v - 1)*Length[v] + Range[Length[v]]
     }], {Length[v], Times @@ Dimensions[a]}];

s.Flatten[a, 1]

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).

Answer 4

This could be fast.

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