使用长度可变的数组列表索引一维数组

时间:2019-08-13 08:39:34

标签: python arrays numpy indexing slice

我想做的事的例子:

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循环,但是它会变慢。

谢谢!

1 个答案:

答案 0 :(得分:2)

我们可以将索引连接起来,用这些索引索引到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中具有较小索引数组的情况而言,这是有意义的。

相关问题