Indexing 1-D array using list of arrays with varying lengths

Question

Example of what I want to do:

import numpy as np

values = np.array([7, 7, 5, 2, 3, 9])
indices = np.array([
    np.array([3,5]), 
    np.array([4]),
    np.array([1,2,3])
    ])

>>> values[indices]
array([
    array([2,9]), 
    array([3]),
    array([7,5,2]),
    ])

Is it possible to achieve this using vectorization? Right now I'm doing it with a for loop, but it can get slow.

Thanks!

Answer 1

We could concatenate the indices, index into values with those and finally split back -

idx = np.concatenate(indices)
all_out = values[idx]
lens = list(map(len,indices))
ssidx = np.r_[0,lens].cumsum()
out = [all_out[i:j] for (i,j) in zip(ssidx[:-1],ssidx[1:])]

For completeness, here's the straight-forward indexing based version -

[values[i] for i in indices]

So, with the proposed method we are making use of slicing and hence reducing per-iteration workload. As such, alongwith the step to get idx that needs concatenation of all indices in the proposed one, it makes sense for the case with small indexing arrays in indices .