时间戳

首页 > 解决方案 > 如何在 numpy 中查找 A 中的行索引，其中第 1 列和第 3 列等于 B（没有 for 循环）？

python - 如何在 numpy 中查找 A 中的行索引，其中第 1 列和第 3 列等于 B（没有 for 循环）？

问题描述

给定 2 个 numpy 数组 A 和 B，如下所示：

A =
[[  1, 200,  15],
 [  0, 600,  25],
 [  2, 200,  20],
 [  3, 100,  10],
 [  1, 300,   5],
 [  4, 400,   3],
 [  0, 100,  12],
 [  5, 300,   2],
 [  2, 300,  25],
 [  6, 100,   1],
 [  1, 400,  74],
 [  7, 300,  10],
 [  2, 400,  15],
 [  8, 200,  13],
 [  3, 400,  29],
 [  3, 600,  37],
 [  0, 200,  20],
 [  9, 300,  42],
 [  5, 400,  30],
 [  6, 200,  51],
 [  7, 400,  21],
 [  9, 500,   5],
 [ 10, 300,   9] ]


B =
[[ 0, 25.],
 [ 1, 74.],
 [ 2, 25.],
 [ 3, 37.],
 [ 4,  3.],
 [ 5, 30.],
 [ 6, 51.],
 [ 7, 21.],
 [ 8, 13.],
 [ 9, 42.],
 [10,  9.]]

计算数组 A 中行索引的有效方法是什么（没有 FOR 循环），其中 A 的第 0 列和第 2 列形成数组 B？

预期的答案是我：

i = [ 1,  5,  8, 10, 13, 15, 17, 18, 19, 20, 22]

以便：

A[i,:] == B

提前谢谢你

标签： pythonnumpy

解决方案

您可以执行以下操作：

  1. 首先抓取Awith的第一列和第三列

    >>> A[:,(0,2)]

  2. (A[i], B[j])然后使用broadcasting技巧和比较所有可能的行对np.equal：

    >>> np.equal(A[:,(0,2)][:,None], B[None])

  3. 确保行中的所有元素都匹配使用np.all：

    >>> np.equal(A[:,(0,2)][:,None], B[None]).all(2)

  4. A最后，查看第一个轴上的 argmax，假设的每一行至少有一个匹配项B：

    >>> np.equal(A[:,(0,2)][:,None], B[None]).all(2).argmax(0)

    B例如，如果您删除最后一行，上述行将不起作用。在其中一行上，所有相等条件都是错误的（在 中根本不匹配B）但 argmax 仍将选择第一个False值，即返回索引0... 这不是一个非常通用的解决方案。

    一种更准确的方法是np.nonzero只捕获非零值，这正是我们在这里想要的：

    >>> np.equal(A[:,(0,2)][:,None], B[None]).all(2).nonzero()[0]

推荐阅读