用另一个数组分割numpy数组



我有一个大的一维整数数组,需要去掉切片。这很琐碎,我只需要做a[start:end]。问题是我需要更多的切片。如果开始和结束是数组,则a[start:end]不起作用。For循环可以用于此操作,但我需要它尽可能快(这是一个瓶颈),因此欢迎使用本地numpy解决方案。

为了进一步说明,我有这个:

a = numpy.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], numpy.int16)
start = numpy.array([1, 5, 7], numpy.int16)
end   = numpy.array([2, 10, 9], numpy.int16)

需要以某种方式将其变成这样:

[[1], [5, 6, 7, 8, 9], [7, 8]] 

这可以(几乎?)在纯numpy中使用掩码数组和步幅技巧来完成。首先,我们创建我们的面具:

>>> indices = numpy.arange(a.size)
>>> mask = ~((indices >= start[:,None]) & (indices < end[:,None]))

或者更简单地说:

>>> mask = (indices < start[:,None]) | (indices >= end[:,None])

对于>=为起始值并且<为结束值的那些索引,掩码是False(即未掩码的值)。(None(又名numpy.newaxis)的切片增加了一个新的维度,实现了广播。)现在我们的口罩看起来是这样的:

>>> mask
array([[ True, False,  True,  True,  True,  True,  True,  True,  True,
True,  True,  True],
[ True,  True,  True,  True,  True, False, False, False, False,
False,  True,  True],
[ True,  True,  True,  True,  True,  True,  True, False, False,
True,  True,  True]], dtype=bool)

现在我们必须使用stride_tricks:来拉伸阵列以适应掩模

>>> as_strided = numpy.lib.stride_tricks.as_strided
>>> strided = as_strided(a, mask.shape, (0, a.strides[0]))
>>> strided
array([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11],
[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11],
[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11]], dtype=int16)

这看起来像一个3x12数组,但每一行都指向相同的内存。现在我们可以将它们组合成一个掩码数组:

>>> numpy.ma.array(strided, mask=mask)
masked_array(data =
[[-- 1 -- -- -- -- -- -- -- -- -- --]
[-- -- -- -- -- 5 6 7 8 9 -- --]
[-- -- -- -- -- -- -- 7 8 -- -- --]],
mask =
[[ True False  True  True  True  True  True  True  True  True  True  True]
[ True  True  True  True  True False False False False False  True  True]
[ True  True  True  True  True  True  True False False  True  True  True]],
fill_value = 999999)

这与你要求的不太一样,但它应该表现得相似。

没有numpy方法可以做到这一点。请注意,由于它是不规则的,所以它无论如何都只是数组/切片的列表。然而,我想补充一点,对于所有(二进制)ufuncs,它们几乎都是numpy中的所有函数(或者它们至少是基于它们的),有reduceat方法,它可能会帮助您避免实际创建切片列表,因此,如果切片很小,也可以加快计算:

In [1]: a = np.arange(10)
In [2]: np.add.reduceat(a, [0,4,7]) # add up 0:4, 4:7 and 7:end
Out[2]: array([ 6, 15, 24])
In [3]: np.maximum.reduceat(a, [0,4,7]) # maximum of each of those slices
Out[3]: array([3, 6, 9])
In [4]: w = np.asarray([0,4,7,10]) # 10 for the total length
In [5]: np.add.reduceat(a, w[:-1]).astype(float)/np.diff(w) # equivalent to mean
Out[5]: array([ 1.5,  5. ,  8. ])

编辑:由于您的切片重叠,我将补充这也可以:

# I assume that start is sorted for performance reasons.
reductions = np.column_stack((start, end)).ravel()
sums = np.add.reduceat(a, reductions)[::2]

[::2]在这里通常应该没什么大不了的,因为对于重叠的切片不需要做真正的额外工作。

对于stop==len(a)。必须避免这种情况。如果你只带一个切片,你可以只做reductions = reductions[:-1](如果是最后一个),但否则你只需要在a上附加一个值来欺骗reduceat:

a = np.concatenate((a, [0]))

因为在末尾添加一个值并不重要,因为您无论如何都要处理切片。

这不是一个"纯"numpy解决方案(尽管正如@mgilson的评论所指出的,很难看出不规则输出是如何成为numpy数组的),但:

a = numpy.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], numpy.int16)
start = numpy.array([1, 5, 7], numpy.int16)
end   = numpy.array([2, 10, 9], numpy.int16)
map(lambda range: a[range[0]:range[1]],zip(start,end))

得到你:

[array([1], dtype=int16), array([5, 6, 7, 8, 9], dtype=int16),  array([7, 8], dtype=int16)]

根据需要。

如果你想把它放在一行,它应该是:

x=[list(a[s:e]) for (s,e) in zip(start,end)]

类似于timday的解决方案。类似速度:

a = np.random.randint(0,20,1e6)
start = np.random.randint(0,20,1e4)
end = np.random.randint(0,20,1e4)
def my_fun(arr,start,end):
return arr[start:end]
%timeit [my_fun(a,i[0],i[1]) for i in zip(start,end)]
%timeit map(lambda range: a[range[0]:range[1]],zip(start,end))

100 loops, best of 3: 7.06 ms per loop 100 loops, best of 3: 6.87 ms per loop

最新更新