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))
How to split an array according to a condition in numpy?