Welcome to OGeek Q&A Community for programmer and developer-Open, Learning and Share
Welcome To Ask or Share your Answers For Others

Categories

0 votes
299 views
in Technique[技术] by (71.8m points)

python - 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 instead of traverse it? Thanks.

See Question&Answers more detail:os

与恶龙缠斗过久,自身亦成为恶龙;凝视深渊过久,深渊将回以凝视…
Welcome To Ask or Share your Answers For Others

1 Reply

0 votes
by (71.8m points)
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]])]

与恶龙缠斗过久,自身亦成为恶龙;凝视深渊过久,深渊将回以凝视…
OGeek|极客中国-欢迎来到极客的世界,一个免费开放的程序员编程交流平台!开放,进步,分享!让技术改变生活,让极客改变未来! Welcome to OGeek Q&A Community for programmer and developer-Open, Learning and Share
Click Here to Ask a Question

...