根据 bool 分割 numpy 数组

Question

我有一个像这样的数组

a = np.array[ 4, 9, 3, 1, 6, 4, 7, 4, 2]

和一个大小相同的布尔数组（所以这是一个掩码）

boo = np.array[ True, True, False, False, True, True, True, False, True]

（

boo

也可以以

False

开头作为第一个条目...）

现在我想将

a

拆分为具有 2 个条件的新数组：

• 新的子数组仅包含

  True

   中带有 

  boo
的值
• 新的子数组始终在

  False

  之后开始，并在

  False

  之前结束。
  所以结果将是

  [[4, 9], [6, 4, 7], [2]]

我的想法是：
我知道我可以使用

np.split

作为基础。
在这种情况下，它将是

b = np.split(a, [2, 4, 7, 8]

，然后我只会从

b

中取出第二个元素，从第一个开始，因为

boo

中的第一个元素是

True

。

所以我的问题是：如何获取数组

[2, 4, 7, 8]

？

（用 python 循环不是一个选择，因为它太慢了。）

Answer 1

也许这已经足够快了：

d = np.nonzero(boo != np.roll(boo, 1))[0]
if d[0] == 0:
    d = d[1:]
b = np.split(a, d)
b = b[0::2] if boo[0] else b[1::2]

---

找到了更简单快捷的方法：

indices = np.nonzero(boo[1:] != boo[:-1])[0] + 1
b = np.split(a, indices)
b = b[0::2] if boo[0] else b[1::2]

比较切片的速度至少是

np.roll()

加 if 语句的两倍。
另外，

np.flatnonzero(...)

看起来会比

np.nonzero(...)[0]

更好，但速度稍慢。

Answer 2

总结最佳解决方案：

def splitByMask(a, mask):
    return numpy.split(a, numpy.nonzero(numpy.diff(mask))[0] + 1)

结果：

>>> myArray = numpy.array([4, 9, 3, 1, 6, 4, 7, 4, 2])
>>> myMask  = numpy.array([1, 1, 0, 0, 1, 1, 1, 0, 1]) # or True / False
>>> splitByMask(myArray, myMask)
[array([4, 9]), array([3, 1]), array([6, 4, 7]), array([4]), array([2])]

这也适用于零长度数组和布尔掩码。