Indexing a list of list of vectors with a list of indices

Question

Given a list of lists of vectors represented as a 3d-array, and a 1d list of indices, how can I index the list of lists such that I return a list of vectors. Semantically, I'd like to achieve the following with a single NumPy call:

N, L, H = 5, 3, 2
data = np.arange(N * L * H).reshape(N, L, H)
inds = np.arange(N) % L
    
indexed_data = []
for x, i in zip(data, inds):
    indexed_data.append(x[i])
y = np.array(indexed_data)

assert y.shape == (N, H)

Seems like np.take should be able to achieve this.

Answer 1

Your example:

In [14]: N, L, H = 5, 3, 2
    ...: data = np.arange(N * L * H).reshape(N, L, H)
    ...: inds = np.arange(N) % L
    ...: 
    ...: indexed_data = []
    ...: for x, i in zip(data, inds):
    ...:     indexed_data.append(x[i])
    ...: y = np.array(indexed_data)

In [15]: y
Out[15]: 
array([[ 0,  1],
       [ 8,  9],
       [16, 17],
       [18, 19],
       [26, 27]])

We can do this in one step by using a row index that matches inds in size:

In [18]: data[np.arange(len(inds)), inds]
Out[18]: 
array([[ 0,  1],
       [ 8,  9],
       [16, 17],
       [18, 19],
       [26, 27]])

The row index:

In [19]: np.arange(len(inds))
Out[19]: array([0, 1, 2, 3, 4])
In [20]: np.r_[:data.shape[0]]
Out[20]: array([0, 1, 2, 3, 4])