按照bool拆分numpy数组

时间:2016-04-09 13:48:18

标签: python arrays numpy split

我有一个类似

的数组

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拆分为具有两个条件的新数组:

  • 新的子数组仅包含True
  • boo的值
  • 新的子数组始终在False后开始,在False之前结束。
    因此结果将是[[4, 9], [6, 4, 7], [2]]

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

所以我的问题是:如何获得数组True

(使用python循环不是一个选项,因为它太慢了。)

1 个答案:

答案 0 :(得分:4)

也许这足够快:

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]更好,但稍微慢一点。