Как разбить массив в соответствии с условием в numpy?
Например, у меня есть ndarray
, который:
a = np.array([1, 3, 5, 7, 2, 4, 6, 8])
Теперь я хочу разделить a
на две части: одна - все числа < 5, а другая - все >= 5:
[array([1,3,2,4]), array([5,7,6,8])]
Конечно, я могу пройти a
и создать два новых массива. Но я хочу знать, имеет ли numpy несколько лучших способов?
Аналогично, для многомерного массива, например
array([[1, 2, 3],
[4, 5, 6],
[7, 8, 9],
[2, 4, 7]])
Я хочу разбить его в соответствии с первым столбцом < 3 и >= 3, результатом которого является:
[array([[1, 2, 3],
[2, 4, 7]]),
array([[4, 5, 6],
[7, 8, 9]])]
Есть ли какие-то лучшие способы вместо того, чтобы пересекать его? Спасибо.
Ответы
Ответ 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)
Это приводит к следующему выводу:
[array([1, 3, 2, 4]), array([5, 7, 6, 8])]
[array([[1, 2, 3],
[2, 4, 7]]), array([[4, 5, 6],
[7, 8, 9]])]