numpy shorthand for taking jagged slice

Question

I have an operation that I\'m doing commonly which I\'m calling a \"jagged-slice\" because I don\'t know the real name for it. It\'s best explained by example:

Answer 1

This is combersome only in the sense that it requires more typing for a task that seems so simple to you.

a[np.arange(a.shape[0]), entries_of_interest]

But as you note, the syntactically simpler a[:, entries_of_interest] has another interpretation in numpy. Choosing a subset of the columns of an array is a more common task that choosing one (random) item from each row.

Your case is just a specialized instance of

a[I, J]

where I and J are 2 arrays of the same shape. In the general case entries_of_interest could be smaller than a.shape[0] (not all the rows), or larger (several items from some rows), or even be 2d. It could even select certain elements repeatedly.

I have found in other SO questions that performing this kind of element selection is faster when applied to a.flat. But that requires some math to construct the I*n+J kind of flat index.

With your special knowledge of J, constructing I seems extra work, but numpy can't make that kind of assumption. If this selection was more common someone could write a function that wraps your expression

def  peter_selection(a,I):
   # check the a.shape[0]==I.shape[0]
   return a[np.arange(a.shape[0]), I]