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.