How to split an array according to a condition in numpy?

Question

For example, I have a ndarray that is:

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

Now I want to split a into two parts, one is all numbers <5 and the other is all >=5:

[array([1,3,2,4]), array([5,7,6,8])]

Certainly I can traverse a and create two new array. But I want to know does numpy provide some better ways?

Similarly, for multidimensional array, e.g.

array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9],
       [2, 4, 7]])

I want to split it according to the first column <3 and >=3, which result is:

[array([[1, 2, 3],
        [2, 4, 7]]), 
 array([[4, 5, 6],
        [7, 8, 9]])]

Are there any better ways instead of traverse it? Thanks.

Asked By: Clippit

||

Source

Answer 1

import numpy as np

def split(arr, cond):
  return [arr[cond], arr[~cond]]

a = np.array([1,3,5,7,2,4,6,8])
print split(a, a<5)

a = np.array([[1,2,3],[4,5,6],[7,8,9],[2,4,7]])
print split(a, a[:,0]<3)

This produces the following output:

[array([1, 3, 2, 4]), array([5, 7, 6, 8])]

[array([[1, 2, 3],
       [2, 4, 7]]), array([[4, 5, 6],
       [7, 8, 9]])]

Answered By: NPE

Answer 2

It might be a quick solution

a = np.array([1,3,5,7])
b = a >= 3 # variable with condition
a[b] # to slice the array
len(a[b]) # count the elements in sliced array

Answered By: Vladimir Gavrysh

Answer 3

1d array
a = numpy.array([2,3,4,…])
a_new = a[(a < 4)] # to get elements less than 5

2d array based on column(consider value of column i should be less than 5,
a = numpy.array([[1,2],[5,6],…]
a = a[(a[:,i] < 5)]

if your condition is multicolumn based, then you can make a new array applying the conditions on the columns. Then you can just compare the new array with value 5(according to my assumption) to get indexes and follow above codes.
Note that, whatever i have written in (), returns the index array.

Answered By: Apurba A

How to split an array according to a condition in numpy?

Question:

Answers: