Question

我有一个形状为[N，M，2]的数组。我对最后一个索引有条件，即

condition = matrix[:,:,0]**2 + matrix[:,:,1]**2 <1

现在，我想选择满足此条件的矩阵元素。如果我盲目地应用matrix[condition]，显然会得到形状为[something，2]的结果，因为否则它将不是矩阵。

我的问题是我需要保持形状 [N，...，2]。例如，最好只获取满足条件的前K个元素（假设始终至少有K个元素）。我该怎么办？

使用for循环，我会写

matrix_new = np.zeros([N,K,2])
for i in range(N):
    condition = matrix[i,:,0]**2 + matrix[i,:,1]**2 <1
    matrix_new[i] = matrix[condition][:K]

但这会非常慢。如何避免for循环？

Answer 1

您已经开拓了新的空间。您只需要将条件矩阵保存到其中：

matrix           = np.array([[[1,2],[3,4]],[[3,4],[5,6]],[[7,8],[0,1]]])
cond             = matrix[...,0]**2+matrix[...,1]**2 < 10 
matrix_new       = np.zeros([3,2,2])
matrix_new[cond] = matrix[cond]

请记住，array [cond]返回numpy视图，然后返回结果。因此，numpy有很多技巧，可以返回您保存的内容，然后返回简化数组。：-）

编辑：

我唯一能找到的就是np.take_along_axis和np.argsort方法，用于基于条件矩阵以特定顺序获取特定向量。像这样：

matrix     = np.array([[[1,2],[3,4]],[[1,1],[5,6]],[[7,8],[0,1]]])
cond       = matrix[...,0]**2 + matrix[...,1]**2
matrix_new = np.take_along_axis(matrix,cond.argsort()[...,None],axis=1)

matrix_new[:,0,:] #Smallest vectors in each. 
matrix_new[:,1,:] #Second smallest vectors in each.

此外，还可以使用“ K = max（（cond <10）.sum（1））”获得值K。

选择满足2D数组中条件的第一个元素

1 个答案: