矢量化 numpy 函数，用于获取其他数组定义的边界之间的数组元素

Question

如果我有

data = np.array([97, 98, 99, 100])
offsets = np.array([0, 1, 2, 2, 3, 3, 4])

那我想做

result = []
for i in range(len(offsets)-1):
    result.append(data[offsets[i]: offsets[i+1]])

并得到

[array([97]),
 array([98]),
 array([], dtype=int64),
 array([99]),
 array([], dtype=int64),
 array([100])]

这可行，但当

data

变大时速度很慢。

我可以在这里使用任何矢量化操作来加快速度吗？

Answer 1

采用多个切片可以转换为高级索引 - 有可能但没有承诺的加速。

第一次尝试你的例子：

In [179]: idx=np.arange(len(data))[:,None]    
In [180]: idx2=(idx>=offsets[:-1])&(idx<offsets[1:])    
In [181]: idx2
Out[181]: 
array([[ True, False, False, False, False, False],
       [False,  True, False, False, False, False],
       [False, False, False,  True, False, False],
       [False, False, False, False, False,  True]])

因此，在每一列中，我们都有

True

来表示我们想要的元素。

使用该掩码的迭代方法是：

In [182]: [data[i] for i in idx2.T]
Out[182]: 
[array([97]),
 array([98]),
 array([], dtype=int32),
 array([99]),
 array([], dtype=int32),
 array([100])]

与你的切片结果相同，但我预计它会更慢。

在其他情况下，像这样的面罩可以压平，只需一步

data[idx]

即可使用。

我在这里能做的最好的事情就是使用

any

，但这会丢失大部分迭代结构

In [183]: idx2.any(axis=1)
Out[183]: array([ True,  True,  True,  True])
In [184]: data[idx2.any(axis=1)]
Out[184]: array([ 97,  98,  99, 100]) 

这

array_split

产生相同的列表，但在内部它也进行迭代/追加

In [221]: np.array_split(data,(1,2,2,3,3))
Out[221]: 
[array([97]),
 array([98]),
 array([], dtype=int32),
 array([99]),
 array([], dtype=int32),
 array([100])]