如何根据其中一个条件有效地分割一对数组

Question

我有两个2D numpy数组 - 真实r，它包含由它们的笛卡尔坐标给出的空间点和v，在每个数组中定义一个复数向量这些点。我想根据r上的一些条件拆分这两个数组。

例如，r1包含具有第一笛卡尔坐标的所有点是正的，并且v1给出v的对应值。所有其他点及其相应的矢量进入。

基于this question，而事实上，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

您可以使用布尔数组作为指标筛选出的值：

先创建一些随机的测试数据：

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的阵列的情况下，我通常使用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直接