Numpy: split array based on groups that satisfy condition

Question

I'm trying to split a numpy array into chunks that obey a certain condition. For example, given a = np.array([ 1, 2, 3, 4, 10, 9, 8, 7, 7, 7, 8, 9, 10]) return [array([3, 4]), array([7, 7, 7])] since those chunks obey the condition np.logical_and(a>2,a<8).

The solution I came up with works, but it feels awkward since I'm applying the condition first to split the array and then again to select the results I want.

import numpy as np
a = np.array([ 1,  2,  3,  4, 10,  9,  8,  7,  7,  7,  8,  9, 10])
cond = np.logical_and(a>2,a<8)
splits = np.split(a,np.where(np.diff(cond))[0]+1)
# this yields [array([1, 2]), array([3, 4]), array([10,  9,  8]), array([7, 7, 7]), array([ 8,  9, 10])]
result = [x for x in splits if np.logical_and(x>2,x<8).any()]
# [array([3, 4]), array([7, 7, 7])]

Is there a way to apply the same condition to both so that if I wanted to change the condition I could change it once in my code instead of twice? Or am I over-complicating things where there's a much simpler, more pythonic, way of doing this?

akuiper · Accepted Answer · 2018-05-14 02:17:41Z

6

Here is one option, subset the array before split to eliminate the for loop:

cond = (a > 2) & (a < 8)
np.split(a[cond], np.where(np.diff(np.where(cond)[0]) > 1)[0] + 1)
# [array([3, 4]), array([7, 7, 7])]

answered May 14, 2018 at 2:17

akuiper

216k33 gold badges362 silver badges379 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

JeffP Over a year ago

That's clever. I forgot that np.where returns an array of indices exactly the same size as the subset so that the original breaks are preserved in the subset by jumps in the indices greater than 1. Thanks!

Collectives™ on Stack Overflow

Numpy: split array based on groups that satisfy condition

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related