python - 在 numpy 数组中查找多个值的行索引

Question

我有一个数组 X:

X = np.array([[4,  2],
              [9,  3],
              [8,  5],
              [3,  3],
              [5,  6]])

我希望在这个数组中找到几个值的行的索引:

searched_values = np.array([[4, 2],
                            [3, 3],
                            [5, 6]])

对于此示例，我希望得到如下结果:

[0,3,4]

我有一个这样做的代码，但我认为它过于复杂:

X = np.array([[4,  2],
              [9,  3],
              [8,  5],
              [3,  3],
              [5,  6]])

searched_values = np.array([[4, 2],
                            [3, 3],
                            [5, 6]])

result = []

for s in searched_values:
    idx = np.argwhere([np.all((X-s)==0, axis=1)])[0][1]
    result.append(idx)

print(result)

我找到了 this answer对于类似的问题，但它仅适用于一维数组。

有没有办法以更简单的方式做我想做的事？

Answer 1

方法#1
一种方法是使用 NumPy broadcasting ，像这样——

np.where((X==searched_values[:,None]).all(-1))[1]

方法#2
一种内存高效的方法是将每一行转换为线性索引等价物，然后使用 np.in1d ，像这样——

dims = X.max(0)+1
out = np.where(np.in1d(np.ravel_multi_index(X.T,dims),\
                    np.ravel_multi_index(searched_values.T,dims)))[0]

方法#3
另一种使用 np.searchsorted 的内存高效方法并且具有相同的转换为线性索引等价物的哲学就像这样 -

dims = X.max(0)+1
X1D = np.ravel_multi_index(X.T,dims)
searched_valuesID = np.ravel_multi_index(searched_values.T,dims)
sidx = X1D.argsort()
out = sidx[np.searchsorted(X1D,searched_valuesID,sorter=sidx)]

请注意这个np.searchsorted方法假设来自 searched_values 的每一行都有一个匹配项在 X .

怎么样 np.ravel_multi_index 工作？
该函数为我们提供了线性索引等效数。它接受 2D n-dimensional indices 的数组，设置为列和 n 维网格本身的形状，这些索引要映射到这些网格上，并且要计算等效的线性索引。
让我们使用我们手头的问题的输入。以输入X为例并注意它的第一行。因为，我们试图转换 X 的每一行转化为它的线性指数等价物，并且自 np.ravel_multi_index假设每一列作为一个索引元组，我们需要转置 X在进入函数之前。因为，X 中每行的元素数在这种情况下是 2 ，要映射到的 n 维网格将为 2D .在 X 中每行 3 个元素，本来是 3D用于映射等的网格。
要了解此函数如何计算线性索引，请考虑 X 的第一行——

In [77]: X
Out[77]: 
array([[4, 2],
       [9, 3],
       [8, 5],
       [3, 3],
       [5, 6]])

我们将 n 维网格的形状设为 dims ——

In [78]: dims
Out[78]: array([10,  7])

让我们创建二维网格，看看该映射是如何工作的，以及如何使用 np.ravel_multi_index 计算线性索引。 ——

In [79]: out = np.zeros(dims,dtype=int)

In [80]: out
Out[80]: 
array([[0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0]])

让我们设置来自 X 的第一个索引元组，即来自 X 的第一行进入网格——

In [81]: out[4,2] = 1

In [82]: out
Out[82]: 
array([[0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 1, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0]])

现在，要查看与刚刚设置的元素等效的线性索引，让我们展平并使用 np.where检测 1 .

In [83]: np.where(out.ravel())[0]
Out[83]: array([30])

如果考虑行优先排序，这也可以计算。
让我们使用 np.ravel_multi_index并验证那些线性指数 -

In [84]: np.ravel_multi_index(X.T,dims)
Out[84]: array([30, 66, 61, 24, 41])

因此，我们将有与来自 X 的每个索引元组相对应的线性索引。 ，即来自 X 的每一行.
为 np.ravel_multi_index 选择尺寸形成唯一的线性索引
现在，考虑每一行 X 背后的想法作为 n 维网格的索引元组并将每个这样的元组转换为标量是具有与唯一元组对应的唯一标量，即 X 中的唯一行.
再来看看X ——

In [77]: X
Out[77]: 
array([[4, 2],
       [9, 3],
       [8, 5],
       [3, 3],
       [5, 6]])

现在，如上一节所述，我们将每一行视为索引元组。在每个这样的索引元组中，第一个元素将代表 n-dim 网格的第一个轴，第二个元素将是网格的第二个轴，依此类推，直到 X 中每行的最后一个元素。 .本质上，每一列将代表网格的一个维度或轴。如果我们要映射来自 X 的所有元素在同一个 n-dim 网格上，我们需要考虑这种建议的 n-dim 网格的每个轴的最大拉伸(stretch)。假设我们正在处理 X 中的正数，这样的拉伸(stretch)将是 X 中每列的最大值+ 1. 那 + 1是因为Python遵循0-based索引。因此，例如 X[1,0] == 9将映射到第 10 行 建议的网格。同样， X[4,1] == 6会去7th栏目 那个网格。
因此，对于我们的示例案例，我们有 -

In [7]: dims = X.max(axis=0) + 1 # Or simply X.max(0) + 1

In [8]: dims
Out[8]: array([10,  7])

因此，我们至少需要一个形状为 (10,7) 的网格。对于我们的示例案例。维度上的更多长度不会受到伤害，并且也会为我们提供独特的线性索引。
结束语:这里要注意的一件重要事情是，如果我们在 X 中有负数，我们需要沿着 X 中的每一列添加适当的偏移量在使用 np.ravel_multi_index 之前将这些索引元组设为正数.