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


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



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}


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


4 Answers 4


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.


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


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.

Diagonal@a[[v]]; // AbsoluteTiming
  {0.420822, Null}

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";

Diagonal@unpacked[[v]]; // AbsoluteTiming
  {0.001294, Null}

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

First@AbsoluteTiming@ Do[Extract[unpacked, Transpose[{v, Range@Length@v}]], {100}] /100
  • $\begingroup$ This is just the sort of thing I was looking for. Funny how easy it seems now but so do many great ideas. Thanks $\endgroup$
    – Mr.Wizard
    Jul 26, 2014 at 17:37

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}]


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.


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.


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

  • 2
    $\begingroup$ If unpacked is not a packed array, Diagonal@unpacked[[v]] is pretty fast and uses almost no extra memory. $\endgroup$
    – Michael E2
    Jul 26, 2014 at 17:25
  • $\begingroup$ @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? $\endgroup$
    – Mr.Wizard
    Jul 26, 2014 at 17:40
  • $\begingroup$ @MichaelE2 Please see Update $\endgroup$
    – Mr.Wizard
    Jul 27, 2014 at 12:03
  • $\begingroup$ 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. :) $\endgroup$
    – Michael E2
    Jul 27, 2014 at 13:10

You could do it with a sparse array:

s = SparseArray[
   MapThread[({#1, #2} -> 1) &,
     (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).

  • $\begingroup$ 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 $\endgroup$
    – Mr.Wizard
    Jul 27, 2014 at 12:10
  • $\begingroup$ 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. $\endgroup$
    – Mr.Wizard
    Jul 27, 2014 at 12:49

This could be fast.

  a[[v[[#]], #]] & /@ Range[Length@v];
  • $\begingroup$ 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? $\endgroup$
    – Mr.Wizard
    Jul 26, 2014 at 17:45
  • $\begingroup$ I checked it with this: SeedRandom[0] a = RandomReal[100, {13000, 13000}]; v = RandomInteger[{1, 13000}, 13000];. I got zero Timing. $\endgroup$ Jul 26, 2014 at 18:19
  • $\begingroup$ For me Table still performs better with that data but this is not far off and it's another option. +1 $\endgroup$
    – Mr.Wizard
    Jul 27, 2014 at 12:07

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.