在numpy数组中搜索

Question

我有一个 2D numpy 数组，比如 A 相对于第 0 列排序。例如

Col.0	Col.1	Col.2
10	2.45	3.25
11	2.95	4
12	3.45	4.25
15	3.95	5
18	4.45	5.25
21	4.95	6
23	5.45	6.25
27	5.95	7
29	6.45	7.25
32	6.95	8
35	7.45	8.25

每行中的条目是唯一的，即第 0 列是 xy 平面中坐标的标识号，第 1 列和第 2 列是这些点的 x 和 y 坐标。 我有另一个数组 B（行可以包含重复数据）。第 0 列和第 1 列存储 x 和 y 坐标。

Col.0	Col.1
2.45	3.25
4.45	5.25
6.45	7.25
2.45	3.25

我的目标是在不使用 for 循环的情况下找到数组 A 中与数组 B 中的数据对应的行索引号。所以，在这种情况下，我的输出应该是[0,4,8,0]。 现在，我知道使用 numpy searchsorted 查找多个数据可以一次性完成。但是，它可以用来与 A 的单列而不是多列进行比较。有没有办法做到这一点？

Answer 1

您可以在 pandas 中使用合并：

df2.merge(df1.reset_index(),how='left',left_on=['Col.0','Col.1'],right_on=['Col.1','Col.2'])['index']

输出：

0    0
1    4
2    8
3    0
Name: index, dtype: int64

如果你喜欢它作为数组：

df2.merge(df1.reset_index(),how='left',left_on=['Col.0','Col.1'],right_on=['Col.1','Col.2'])['index'].to_numpy()
#array([0, 4, 8, 0])

Answer 2

纯 numpy 解决方案：

我的直觉是，我通过广播获取a[:,1:] 和b 之间的c 差异，因此c 的形状为(11, 4, 2)。匹配的行将全为零。然后我做c == False 来获得一个面具。我做c.all(2) 会产生一个形状为(11, 4) 的布尔数组，其中所有True 元素表示a 和b 之间的匹配。然后我只需使用np.nonzero 来获取所述元素的索引。

import numpy as np

a = np.array([
    [10, 2.45, 3.25],
    [11, 2.95, 4],
    [12, 3.45, 4.25],
    [15, 3.95, 5],
    [18, 4.45, 5.25],
    [21, 4.95, 6],
    [23, 5.45, 6.25],
    [27, 5.95, 7],
    [29, 6.45, 7.25],
    [32, 6.95, 8],
    [35, 7.45, 8.25],
])

b = np.array([
    [2.45, 3.25],
    [4.45, 5.25],
    [6.45, 7.25],
    [2.45, 3.25],
])

c = (a[:,np.newaxis,1:]-b) == False
rows, cols = c.all(2).nonzero()
print(rows[cols.argsort()])
# [0 4 8 0]