python - 根据条件获取 NumPy 数组的连续元素组

Question

我有一个 NumPy 数组如下:

import numpy as np
a = np.array([1, 4, 2, 6, 4, 4, 6, 2, 7, 6, 2, 8, 9, 3, 6, 3, 4, 4, 5, 8])

和一个常数b = 6

基于 previous question我可以计算 c 的数量，它由 a 中的元素连续 2 次或更多次小于 b 的次数定义。

from itertools import groupby
b = 6
sum(len(list(g))>=2 for i, g in groupby(a < b) if i)

所以在这个例子中 c == 3

现在我想在每次满足条件时输出一个数组，而不是计算条件满足的次数。

所以对于这个例子，正确的输出是:

array1 = [1, 4, 2]
array2 = [4, 4]
array3 = [3, 4, 4, 5]

自:

1, 4, 2, 6, 4, 4, 6, 2, 7, 6, 2, 8, 9, 3, 6, 3, 4, 4, 5, 8  # numbers in a
1, 1, 1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 0, 1, 0, 1, 1, 1, 1, 0  # (a<b)
^^^^^^^-----^^^^-----------------------------^^^^^^^^^^---  # (a<b) 2+ times consecutively
   1         2                                    3

到目前为止，我尝试了不同的选择:

np.isin((len(list(g))>=2 for i, g in groupby(a < b)if i), a)

和

np.extract((len(list(g))>=2 for i, g in groupby(a < b)if i), a)

但他们都没有实现我正在寻找的东西。有人可以为我指出正确的 Python 工具以输出满足我的条件的不同数组吗？

Answer 1

同时测量 my other answer 的性能我注意到虽然它比 Austin's solution 快(对于长度 <15000 的数组)，其复杂度不是线性的。

基于 this answer我使用 np.split 想出了以下解决方案这比之前在此处添加的两个答案都更有效:

array = np.append(a, -np.inf)  # padding so we don't lose last element
mask = array >= 6  # values to be removed
split_indices = np.where(mask)[0]
for subarray in np.split(array, split_indices + 1):
    if len(subarray) > 2:
        print(subarray[:-1])

给出:

[1. 4. 2.]
[4. 4.]
[3. 4. 4. 5.]

性能*:

^{*由 perfplot 测量}