Как эффективно разбить пару массивов на основе условия на одном из них

Question

У меня есть два двумерных массива - вещественный r, который содержит точки в пространстве, заданные их декартовыми координатами, и v, комплексный вектор, определенный в каждой из этих точек. Я хотел бы разделить оба этих массива, основываясь на некотором условии на r.

например, r1 содержит все точки с положительной первой декартовой координатой, а v1 дает соответствующие значения v. Все остальные точки и соответствующие им векторы входят в.

Исходя из этого вопроса и того факта, что zip по сути своей является обратным, у меня в настоящее время есть следующее решение:

r1, v1 = zip(*[rv for rv in zip(r, v) if rv[0][0] > 0.0])
r2, v2 = zip(*[rv for rv in zip(r, v) if rv[0][0] <= 0.0])
r1 = np.array(r1)
r2 = np.array(r2)
v1 = np.array(v1)
v2 = np.array(v2)

Это работает достаточно хорошо для моих целей, однако включает преобразование в большие списки массивов, что, безусловно, весьма неэффективно.

Существует ли альтернативное решение, которое является быстрым, лаконичным и позволяет избежать создания промежуточных списков?

Answer 1

вы можете использовать массив bool в качестве индекса для фильтрации значений:

сначала создайте несколько случайных тестовых данных:

import numpy as np
np.random.seed(0)
r = np.random.rand(10,2)-0.5
v = np.random.rand(10) + np.random.rand(10)*1j

, то:

idx = r[:,0] > 0 # idx is a bool array 
r1 = r[idx]
v1 = v[idx]

r2 = r[~idx] # ~idx compute bit-wise NOT, element-wise
v2 = v[~idx]

Answer 2

При проверке условий числовых массивов я обычно заканчиваю с использованием numpy.where, только с условием в качестве аргументов, он возвращает индексы массива:

i1 = numpy.where(r[:,0]>0.0) # i1 has now the row indices where column 0 > 0.0
i2 = numpy.where(r[:,0]<=0.0)
r1 = numpy.take(r,i1[0],0)       # take slices of r along axis 0
v1 = numpy.take(v,i1[0],0)
r2 = numpy.take(r,i2[0],0)
v2 = numpy.take(v,i2[0],0)

Несколько короче, просто используйте сжатиев этом случае, который сочетает в себе:

larger = r[:,0]>0.0
r1 = numpy.compress(larger,r,0)

Я не знаю, если это быстрее, но он использует только массивы, без промежуточных списков

РЕДАКТИРОВАТЬ: вы также можете посмотреть назамаскированные массивы, если вы хотите работать с r, v напрямую