python - 如何在不迭代的情况下基于成对的开始/结束索引定义 numpy 数组的多个切片？

Question

我有一个 numpy 整数数组。

我有两个其他数组表示开始和长度(或者它可以是开始和结束)索引到这个数组中，它们标识我需要处理的整数序列。序列是可变长度的。

x=numpy.array([2,3,5,7,9,12,15,21,27,101, 250]) #Can have length of millions

starts=numpy.array([2,7]) # Can have lengths of thousands
ends=numpy.array([5,9])

# required output is x[2:5],x[7:9] in flat 1D array 
# [5,7,9,12,21,27,101] 

我可以使用 for 循环轻松完成此操作，但该应用程序对性能敏感，因此我正在寻找一种无需 Python 迭代即可完成此操作的方法。

我们将不胜感激!

道格

Answer 1

方法 #1

一种矢量化方法是通过广播创建屏蔽 -

In [16]: r = np.arange(len(x))

In [18]: x[((r>=starts[:,None]) & (r<ends[:,None])).any(0)]
Out[18]: array([ 5,  7,  9, 21, 27])

方法 #2

另一种矢量化方法是使用 cumsum 创建 1 和 0 的斜坡(对于许多开始-结束对应该更好)，就像这样 -

idx = np.zeros(len(x),dtype=int)
idx[starts] = 1
idx[ends[ends<len(x)]] = -1
out = x[idx.cumsum().astype(bool)]

方法 #3

另一个实现内存效率的基于循环的方法，在 starts,ends 对中有很多条目可能会更好 -

mask = np.zeros(len(x),dtype=bool)
for (i,j) in zip(starts,ends):
    mask[i:j] = True
out = x[mask]

方法 #4

为了完整起见，这里有另一个带循环的选择切片，然后分配到一个初始化的数组中，应该适合从大数组中选择的切片 -

lens = ends-starts
out = np.empty(lens.sum(),dtype=x.dtype)
start = 0
for (i,j,l) in zip(starts,ends,lens):
    out[start:start+l] = x[i:j]
    start += l

如果迭代次数很多，可以进行较小的优化以减少每次迭代的计算量 -

lens = ends-starts
lims = np.r_[0,lens].cumsum()
out = np.empty(lims[-1],dtype=x.dtype)
for (i,j,s,t) in zip(starts,ends,lims[:-1],lims[1:]):
    out[s:t] = x[i:j]