在二维数组中选择满足条件的第一个元素答案

【问题标题】：Select first elements that satisfy a condition in a 2D array在二维数组中选择满足条件的第一个元素
【发布时间】：2020-04-27 16:54:55
【问题描述】：

我有一个形状为 [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 循环？

【问题讨论】：

尝试向量化您的代码。看这个例子stackoverflow.com/questions/47369976/…
你问的不清楚。您不能简单地为满足条件的元素设置形状 [N, ..., 2]，因为每行中可能有不同数量的元素满足条件，因此输出不会是矩形，因此 numpy 无法创建数组出它。您的意思是让不满足条件的元素保留一个默认值，比如 0？
@Ehsan 我现在明白你的意思了。是的，我假设 OP 希望将失败替换为零，否则 OP 只需要 matrix[cond] 本身。
不，我不想用零替换。我想排除不满意的项目。为了得到一个矩阵，我想取满足条件的“前K”个元素（假设总是至少有K）
@AmitGupta 但这就是问题所在，我该如何对其进行矢量化？

标签： python numpy

【解决方案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

【讨论】：

这会将不满足条件的保留为零。
我没有遵循你的意思。您还希望他们成为什么？
这不是我的问题，也没有说明它打算对它们做什么。我没有批评你的回答。只需一个评论，以供读者注意。
@Ehsan 也是我的错，没有意识到你不是 OP。
这不应该发生，因为在这种情况下，每行中满足约束的最小元素数是 K=0。这意味着我想要在您的特定情况下使用空矩阵