Select consecutive elements that satisfy a certain condition as separate arrays

Question

Given an array of values, I want to select multiple sequences of consecutive elements that satisfy a condition. The result should be one array for each sequence of elements.

For example I have an array containing both negative and positive numbers. I need to select sequences of negative numbers, with each sequence in a separate array.
Here is an example :

import numpy as np

# Example data
values = np.array([1, 2, 3, -1, -2, -3, 4, 5, 6, -7, -8, 10])

mask = values < 0

Here is how the output should look like :

Array 1:

[-1 -2 -3]

Array 2:

[-7 -8]

I tried to do it using numpy.split, but it became more like spaghetti code. I was wondering is there a Pythonic way to do this task?

Answer 1

clustering groups of negative numbers

If you only want to group the chunks of negative values, irrespective of their relative values, then simply compute a second mask to identify the starts of each negative chunk:

mask = values < 0
mask2 = np.r_[True, np.diff(mask)]
out = np.array_split(values[mask], np.nonzero(mask2[mask])[0][1:])

Output: [array([-1, -7, -3]), array([-7, -8])]

clustering groups of negative numbers if they are successive in value

If you want to cluster the negative values that also a successively decreasing (e.g. -1, -2, -3, -5, -6 would form 2 clusters: -1, -2, -3 and -5, -6. Then I would use pandas:

• convert to Series
• identify the negative values
• create a grouper for consecutive negative values ((~mask).cumsum())
• add the index (or a range) to group the successive
• groupby

import pandas as pd

s = pd.Series(values)

# mask to keep negative values
mask = s<0
# group consecutive negatives
group1 = (~mask).cumsum()
# group successive decrementing values
s2 = s+s.index
group2 = s2.ne(s2.shift()).cumsum()

out = [g.to_numpy() for k, g in s[mask].groupby([group1, group2])]

Output: [array([-1, -2, -3]), array([-7, -8]), array([-7])]

Intermediates:

     s   mask  s2  group1  group2
0    1  False   1       1       1
1    2  False   3       2       2
2    3  False   5       3       3
3   -1   True   2       3       4 # out 1
4   -2   True   2       3       4 #
5   -3   True   2       3       4 #
6    4  False  10       4       5
7    5  False  12       5       6
8    6  False  14       6       7
9   -7   True   2       6       8  # out 2 
10  -8   True   2       6       8  #
11  -7   True   4       6       9    # out 3
12  10  False  22       7      10

Answer 2

You could use itertools.groupby to get the result you want:

import itertools

out = [list(g) for k, g in itertools.groupby(values, key=lambda v: v < 0) if k]

Output:

[[-1, -2, -3], [-7, -8]]