使用具有不同长度的数组列表索引一维数组

Question

我想做的事情的例子：

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]),
])

是否可以使用矢量化来实现这一点？ 现在我正在使用 for 循环来做这件事，但它可能会变慢。

谢谢！

Answer 1

我们可以将索引连接起来，索引与索引一起values，最后拆分回来 -

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:])]

为了完整起见，这里是基于索引的简单版本-

[values[i] for i in indices]

因此，通过所提出的方法，我们正在利用切片，从而减少每次迭代的工作量。因此，随着获取需要连接建议索引中所有索引的idx的步骤，对于indices中的小型索引数组的情况是有意义的。