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

您可以使用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数组的条件时，我通常最终使用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上操作，你可能还想查看掩码数组

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

2 个答案: