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

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 in开发者_开发问答stead of traverse it? Thanks.

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]])]

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

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.