使用NumPy查找条件为True的范围

时间:2013-06-17 15:14:25

标签: python arrays search numpy

想象一下,我有一个numpy数组,我需要找到条件为True的跨度/范围。例如,我有以下数组,其中我试图找到项大于1的跨度:

[0, 0, 0, 2, 2, 0, 2, 2, 2, 0]

我需要找到索引(开始,停止):

(3, 5) 
(6, 9)

我能够实现的最快的事情是制作一个布尔数组:

truth = data > threshold

然后使用numpy.argminnumpy.argmax循环遍历数组以查找开始和结束位置。

    pos = 0
    truth = container[RATIO,:] > threshold

    while pos < len(truth):
        start = numpy.argmax(truth[pos:]) + pos + offset
        end = numpy.argmin(truth[start:]) + start  + offset
        if not truth[start]:#nothing more
            break
        if start == end:#goes to the end
            end = len(truth)
        pos = end

但是对于阵列中的数十亿个位置来说这已经太慢了,而且我发现的跨度通常只是连续几个位置。有没有人知道找到这些跨度的更快方法?

2 个答案:

答案 0 :(得分:5)

如何一种方式。首先取你拥有的布尔数组:

In [11]: a
Out[11]: array([0, 0, 0, 2, 2, 0, 2, 2, 2, 0])

In [12]: a1 = a > 1

使用roll将其向左移一(以获得每个索引的下一个状态):

In [13]: a1_rshifted = np.roll(a1, 1)

In [14]: starts = a1 & ~a1_rshifted  # it's True but the previous isn't

In [15]: ends = ~a1 & a1_rshifted

其中non-zero是每个True批次的开头(或者分别是结束批次):

In [16]: np.nonzero(starts)[0], np.nonzero(ends)[0]
Out[16]: (array([3, 6]), array([5, 9]))

将这些拼凑在一起:

In [17]: zip(np.nonzero(starts)[0], np.nonzero(ends)[0])
Out[17]: [(3, 5), (6, 9)]

答案 1 :(得分:1)

如果您有权访问scipy库:

您可以使用scipy.ndimage.measurements.label来识别任何非零值区域。它返回一个数组,其中每个元素的值是原始数组中范围或范围的id。

然后,您可以使用scipy.ndimage.measurements.find_objects返回提取这些范围所需的切片。您可以直接从这些切片访问开始/结束值。

在你的例子中:

from numpy import array
from scipy.ndimage.measurements import label, find_objects

data = numpy.array([0, 0, 0, 2, 2, 0, 2, 2, 2, 0])

labels, number_of_regions = label(a)
ranges = find_objects(labels)

for identified_range in ranges:
    print identified_range[0].start, identified_range[0].stop

您应该看到:

3 5
6 9

希望这有帮助!