Split numpy array into segments where condition is met

Question

I have an array like so:

arr = np.array([1, 2, 3, 4, -5, -6, 3, 5, 1, -2, 5, -1, -1, 10])

I want to get rid of all negative values, and split the array at each index where there was a negative value. The result should look like this:

split_list = [[1, 2, 3, 4], [3, 5, 1], [5], [10]]

I know how to do this using list comprehension, but since the array can get quite large and I have to do the calculation many times, I want to find a solution using numpy. I found this https://www.geeksforgeeks.org/python-split-list-into-lists-by-particular-value/, which I can use to split the array where there are negative values, but I can't simultaneously remove them.

Answer 1

Note that instead of numpy, you could make use of itertools.groupby this way (though, judging on this: NumPy grouping using itertools.groupby performance, pure numpy will likely be more efficient):

import numpy as np
from itertools import groupby

arr = np.array([1, 2, 3, 4, -5, -6, 3, 5, 1, -2, 5, -1, -1, 10])
split_list = [list(group) for key, group in groupby(arr, key=lambda x:x>=0) if key]

# [[1, 2, 3, 4], [3, 5, 1], [5], [10]]