How to split a numpy array to 2d array based on postive/neagtive changes

Question

I have a numpy 1D array:

import numpy as np

arr = np.array([1,  1,  3, -2, -1,  2,  0,  2,  1,  1, -3, -1,  2])

I want split it into another two-dimensional array, based on changes in positive and negative values of array's elements(0 is placed in the range of positive values). But the original order of elements should be maintained.

The desired result is:

new_arr = [[1, 1, 3], [-2, -1], [2, 0, 2, 1, 1], [-3, -1], [2]]

Answer 1

You could use array_split, diff, nonzero:

np.array_split(arr, np.nonzero(np.diff(arr>=0))[0]+1)

Ouptut:

[array([1, 1, 3]),
 array([-2, -1]),
 array([2, 0, 2, 1, 1]),
 array([-3, -1]),
 array([2])]

Intermediates:

# arr>0
[ True  True  True False False  True False  True  True  True False False  True]

# np.diff(arr>=0)
[False False  True False  True False False False False  True False  True]

# np.nonzero(np.diff(arr>=0))[0]+1
[ 3  5 10 12]

And for lists as output:

out = list(map(list, np.array_split(arr, np.nonzero(np.diff(arr>=0))[0]+1)))

Output:

[[1, 1, 3], [-2, -1], [2, 0, 2, 1, 1], [-3, -1], [2]]

Or using itertools.groupby:

from itertools import groupby

out = [list(g) for _,g in groupby(arr, key=lambda x: x>=0)]

Output:

[[1, 1], [3, -2], [-1, 2, 0, 2, 1], [1, -3], [-1, 2]]

comparison of approaches