Skip to main content
 首页 » 编程设计

python之根据 bool 拆分 numpy 数组

2025年05月04日26arxive

我有一个类似的数组

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

和一个大小相同的 bool 数组(所以这是一个掩码)

boo = np.array[真,真,假,假,真,真,真,假,真]

(boo 也可以以 False 作为第一个条目...)

现在我想将 a 拆分为具有 2 个条件的新数组:

  • 一个新的子数组只包含 boo 中的 True 值
  • 新的子数组总是在 False 之后开始并在 False 之前结束。
    所以结果将是 [[4, 9], [6, 4, 7], [2]]

我的想法是:
我知道我可以使用 np.split作为基本。
在这种情况下,它将是 b = np.split(a, [2, 4, 7, 8] 然后我只会从 b 中获取第二个元素code>,从第一个开始,因为我在 boo 中的第一个元素是 True。

所以我的问题是:如何获取数组 [2, 4, 7, 8]?

(用 python 循环不是一种选择,因为它太慢了。)

请您参考如下方法:

也许这已经足够快了:

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] 更好,但速度稍慢。