从掩码中高效地提取numpy子数组

Question

我正在寻找一种使用掩码从给定数组中提取多个子数组的 Pythonic 方法，如示例所示：

a = np.array([10, 5, 3, 2, 1])
m = np.array([True, True, False, True, True])

输出结果将是一个数组的集合，例如以下内容，其中仅由掩码 m 中连续的“区域”中的 True 值（相邻的 True 值）表示生成子数组的索引。

L[0] = np.array([10, 5])
L[1] = np.array([2, 1])

- TommasoF

---

另一个方法是使用scipy.ndimage.measurements.label，如https://dev59.com/omox5IYBdhLWcg3wFAjW中所建议的。 - TommasoF

Answer 1

这里有一种方法 -

def separate_regions(a, m):
    m0 = np.concatenate(( [False], m, [False] ))
    idx = np.flatnonzero(m0[1:] != m0[:-1])
    return [a[idx[i]:idx[i+1]] for i in range(0,len(idx),2)]

示例运行 -

In [41]: a = np.array([10, 5, 3, 2, 1])
    ...: m = np.array([True, True, False, True, True])
    ...: 

In [42]: separate_regions(a, m)
Out[42]: [array([10,  5]), array([2, 1])]

运行时测试

其他方法 -

# @kazemakase's soln
def zip_split(a, m):
    d = np.diff(m)
    cuts = np.flatnonzero(d) + 1

    asplit = np.split(a, cuts)
    msplit = np.split(m, cuts)

    L = [aseg for aseg, mseg in zip(asplit, msplit) if np.all(mseg)]
    return L

时间 -

In [49]: a = np.random.randint(0,9,(100000))

In [50]: m = np.random.rand(100000)>0.2

# @kazemakase's's solution
In [51]: %timeit zip_split(a,m)
10 loops, best of 3: 114 ms per loop

# @Daniel Forsman's solution
In [52]: %timeit splitByBool(a,m)
10 loops, best of 3: 25.1 ms per loop

# Proposed in this post
In [53]: %timeit separate_regions(a, m)
100 loops, best of 3: 5.01 ms per loop

增加岛屿的平均长度 -

In [58]: a = np.random.randint(0,9,(100000))

In [59]: m = np.random.rand(100000)>0.1

In [60]: %timeit zip_split(a,m)
10 loops, best of 3: 64.3 ms per loop

In [61]: %timeit splitByBool(a,m)
100 loops, best of 3: 14 ms per loop

In [62]: %timeit separate_regions(a, m)
100 loops, best of 3: 2.85 ms per loop

- Divakar

---

我接受这个答案，因为它提供了与其他讨论方法的比较，此外还有一种更快的方法。谢谢！ - TommasoF

有趣的事实：我刚刚发现 np.r_[False, m, False] 比 np.concatenate(([False], m, [False])) 慢5-10倍。 - Mad Physicist

Answer 2

def splitByBool(a, m):
    if m[0]:
        return np.split(a, np.nonzero(np.diff(m))[0] + 1)[::2]
    else:
        return np.split(a, np.nonzero(np.diff(m))[0] + 1)[1::2] 

这将返回一个数组列表，将其分成由m中的True定义的块。

- Daniel F

---

不错的解决方案。利用了“True”和“False”段必须交替出现的事实。 - MB-F

我喜欢这个解决方案，因为它可以转换成一个单行代码：np.split(a, np.nonzero(np.diff(m))[0] + 1)[1 - m[0]::2]。 - Mad Physicist

甚至可以使用np.split(a, np.flatnonzero(np.diff(m)) + 1)[1 - m[0]::2)]，这样更易读一些。 - Daniel F

有趣的是，每个一行代码都比上一个慢 :) 但正如你所说，更容易阅读。 - Mad Physicist

Answer 3

看起来np.split是一个自然的应用程序。

你首先需要确定在哪里切割数组，这是掩码在True和False之间变化的地方。接下来丢弃所有掩码为False的元素。

a = np.array([10, 5, 3, 2, 1])
m = np.array([True, True, False, True, True])

d = np.diff(m)
cuts = np.flatnonzero(d) + 1

asplit = np.split(a, cuts)
msplit = np.split(m, cuts)

L = [aseg for aseg, mseg in zip(asplit, msplit) if np.all(mseg)]

print(L[0])  # [10  5]
print(L[1])  # [2 1]

- MB-F