【提问与解答】使用 numpy 找出两个二维数组相同元素的索引

---

问题描述：

给定两个二维数组，bigArray，smallArray，这两个数组都用来保存一系列的二维坐标点 (x, y)，现在需要将 smallArray 中与 bigArray 中的相同坐标点在 bigArray 中的索引找出来。

由于 smallArray 与 bigArray 中保存的数据众多，因此使用 双 for 循环效率低下。可以使用 numpy 的 == 以及 where 函数进行替代操作。现记录一下相关的语法示例。

---

示例代码1：

import numpy as np

# 创建一个 shape 为 (3, 2) 的二维大数组
testBigArray = np.array([[1, 2],
       				     [3, 4],
       				     [5, 6]])

# 创建一个 shape 为 (1, 2) 的二维小数组
testSmallArray = np.array([[3, 4]])

# numpy 广播比较两数组元素 (3, 2) -> (3, 1, 2) : (行，列）-> （行，1，列）
result_1 = testBigArray[:, None] == testSmallArray
print(result_1)

# 进一步操作，对 result_1 结果沿最后一个轴检查每一行是否为 true
result_2 = (testBigArray[:, None] == testSmallArray).all(-1)
print(result_2)

# 进一步操作，使用 np.where() 找出包含 True 元素的索引
result_3 = np.where((testBigArray[:, None] == testSmallArray).all(-1)) 
print(result_3)

# 进一步操作，提取元素相等的 testBigArray 中的行索引
result_4 = np.where((testBigArray[:, None] == testSmallArray).all(-1))[0]
print(result_4) 

输出 result_1：

""" result_1 的结果
testBigArray 通过 None 的方式在对应维度添加一维，shape 从 (3, 2) 扩充为 (3, 1, 2)，
可以写成 testBigArray[:, None, :]。

testBigArray[:, None][0][0]
Out[115]: array([1, 2])
testBigArray[:, None][1][0]
Out[116]: array([3, 4])
testBigArray[:, None][2][0]
Out[117]: array([5, 6])

testBigArray[:, None] 中的每个 [x, y] 元素与 smallArray 中的元素 [x, y] 作比较
（x对比x，并返回 bool 结果，同理之于y）（？描述是否准确）。这段代码生成一个bool数组，
shape 同 testBigArray[:, None] 为 (3, 1, 2)

"""
array([[[False, False]],
       [[ True,  True]],
       [[False, False]]])

输出 result_2：

""" result_2结果
result_2 会在 生成的 result_1，shape 为 (3, 1, 2) 的 bool 类型数组，依照最后一个维度（2）（列），
逐行进行 ‘与’ 操作，并生成 result_2 shape 为 (3, 1) 的 bool 类型矩阵（相当于降维，列维度没了）。
"""
array([[False],
       [ True],
       [False]])

输出 result_3, result_4：

""" result_3 结果
np.where(condition) 返回一个 length 为 2 的元组，第一个元素表示 condition 生成的 bool 数组中
True 元素所对应的 行索引，第二个元素表示对应的列索引。

查看 result_2 的结果，可以发现 True 元素对应的 行索引为 1，对应的列索引为 0。
"""
Out[124]: (array([1], dtype=int64), array([0], dtype=int64))

""" result_4 结果
"""
array([1], dtype=int64)

---

示例代码2：

使用复杂一点的 testSmallArray。

import numpy as np

# 创建一个 shape 为 (3, 2) 的二维大数组
testBigArray = np.array([[1, 2],
       				     [3, 4],
       				     [5, 6]])

# 创建一个 shape 为 (1, 2) 的二维小数组
testSmallArray3 = np.array([[3, 4],
       						[7, 8],
       						[5, 6]])

# numpy 广播比较两数组元素 (3, 2) -> (3, 1, 2) : (行，列）-> （行，1，列）
result_1 = testBigArray[:, None] == testSmallArray3
print(result_1)

# 进一步操作，对 result_1 结果沿最后一个轴检查每一行是否为 true
result_2 = (testBigArray[:, None] == testSmallArray3).all(-1)
print(result_2)

# 进一步操作，使用 np.where() 找出包含 True 元素的索引
result_3 = np.where((testBigArray[:, None] == testSmallArray3).all(-1)) 
print(result_3)

# 进一步操作，提取元素相等的 testBigArray 中的行索引
result_4 = np.where((testBigArray[:, None] == testSmallArray3).all(-1))[0]
print(result_4) 

输出结果：

""" result_1
testBigArray[:, None]
Out[130]: 
array([
       [[1, 2]],
       [[3, 4]],
       [[5, 6]]
      ])
"""
array([
	   [[False, False],
        [False, False],
        [False, False]],  # testBigArray 中的 [1, 2] 与 smallArray3 中的 [3, 4], [7, 8], [5, 6] 分别对比的结果
       [[ True,  True],
        [False, False],
        [False, False]],  # 同理
       [[False, False],
        [False, False],
        [ True,  True]]   # 同理
      ])

""" result_2
"""
Out[133]: 
array([[False, False, False],
       [ True, False, False],
       [False, False,  True]])

""" result_3
参照 result_2，其有两个 True 元素，他们所在的行索引为 1 和 2，列索引为 0 和 2
"""
Out[134]: (array([1, 2], dtype=int64), array([0, 2], dtype=int64))

""" result_4
因为我们需要行索引，因此提取 result_3 的第一个元素
"""
Out[135]: array([1, 2], dtype=int64)

---

记录 ChatGPT 关于这个问题的解答

问题1：请创建一个int类型 shape 为 (10, 2) 的 bigArray 和一个 int类型 shape 为 (3, 2) 的 smallArray, 请返回bigArray 中与 smallArray 相等的 行 的 index

回答：

可以使用 np.where 函数来查找 bigArray 中与 smallArray 相等的行的索引。代码如下：

import numpy as np

bigArray = np.array([[1, 2],
                     [3, 4],
                     [5, 6],
                     [7, 8],
                     [9, 10],
                     [11, 12],
                     [13, 14],
                     [15, 16],
                     [17, 18],
                     [19, 20]])

smallArray = np.array([[3, 4],
                       [9, 10],
                       [13, 14]])

# 使用 np.where 查找相等的行的索引
index = np.where((bigArray[:, None] == smallArray).all(-1))[0]

print(index)

输出结果为：

[1 4 6]

这个代码中，bigArray[:, None] 将 bigArray 转换为一个形状为 (10, 1, 2) 的三维数组，其中每个元素都是 bigArray 中的一行。smallArray 是一个形状为 (3, 2) 的二维数组。将它与 bigArray 进行比较时，使用了广播（broadcasting）机制，使得 smallArray 的每一行都与 bigArray 中的所有行逐一比较。比较结果是一个形状为 (10, 3) 的二维数组，其中每一行表示 bigArray 中对应位置的元素是否与 smallArray 中的元素相等。.all(-1) 方法沿着最后一个轴（也就是列轴）检查每一行的所有元素是否都为 True。如果某一行的所有元素都为 True，则返回 True；否则返回 False。最后，np.where 函数返回一个包含满足条件的元素的行索引的一维数组。

问题2：请问 bigArray[:, None] == smallArray 的运行逻辑是什么

bigArray[:, None] 是将 bigArray 的维度从 (10, 2) 扩展为 (10, 1, 2) 的操作。这样做是为了与 smallArray 进行比较时进行广播（broadcasting），使得 smallArray 的每一行都可以与 bigArray 中的所有行逐一比较。

bigArray[:, None] 的结果是一个形状为 (10, 1, 2) 的三维数组，其中每个元素都是 bigArray 中的一行。这个操作可以通过在 bigArray 的第二个轴上插入一个新轴来实现。插入新轴后，bigArray 的形状变为 (10, 1, 2)，其中第一个轴表示行，第二个轴表示插入的新轴，第三个轴表示列。

接下来，bigArray[:, None] == smallArray 将 bigArray[:, None] 和 smallArray 进行逐元素比较，生成一个形状为 (10, 3, 2) 的布尔型数组。其中，第一个轴表示 bigArray 中的行，第二个轴表示 smallArray 中的行，第三个轴表示列。比较的结果是一个布尔型数组，其中的每个元素表示 bigArray 中对应位置的元素是否与 smallArray 中的元素相等。

例如，对于 bigArray 中的第一行 [1, 2]，与 smallArray 中的任何一行都不相等，因此该行的结果为 [False, False, False]。对于 bigArray 中的第二行 [3, 4]，与 smallArray 的第一行 [3, 4] 相等，因此该行的结果为 [True, False, False]。以此类推，对 bigArray 中的每一行进行相应的比较。

最后，(bigArray[:, None] == smallArray).all(-1) 使用 .all(-1) 方法沿着最后一个轴（也就是列轴）检查每一行的所有元素是否都为 True。如果某一行的所有元素都为 True，则返回 True；否则返回 False。这样得到的结果是一个形状为 (10, 3) 的布尔型数组，其中的每个元素表示 bigArray 中对应位置的行是否与 smallArray 中的任何一行完全相等。

希望这样解释清楚了 bigArray[:, None] == smallArray 的运行逻辑。如果还有其他问题，请随时提问。

问题3：接上面的问题，np.where 作用是什么

np.where 函数是 NumPy 中的一个函数，用于根据给定条件返回满足条件的元素的索引。

np.where 函数的基本语法是：

np.where(condition, [x, y])

其中，condition 是一个条件表达式，可以是一个布尔型数组或一个布尔型条件表达式。x 和 y 是可选参数，用于指定满足条件和不满足条件时返回的值。如果只提供了 condition 参数，则 np.where 函数返回满足条件的元素的索引。

np.where 函数的返回值是一个包含满足条件的元素的索引的元组。如果 condition 是一个布尔型数组，返回的索引是一个包含满足条件的元素的位置的一维数组。如果 condition 是一个布尔型条件表达式，返回的索引是一个包含满足条件的元素的位置的元组，其中每个元素都是一个一维数组。

在上面的代码中，np.where((bigArray[:, None] == smallArray).all(-1)) 返回满足条件 (bigArray[:, None] == smallArray).all(-1) 的元素的索引。这个条件是一个布尔型数组，表示 bigArray 中的每一行是否与 smallArray 中的任何一行完全相等。np.where 函数返回一个一维数组，其中包含满足条件的元素的索引。最后，通过 [0] 取出这个一维数组的第一个元素，得到一个包含满足条件的元素的索引的一维数组。

希望这样解释清楚了 np.where 函数的作用。如果还有其他问题，请随时提问。

---

参考：

链接1：Numpy数组中None的作用
链接2：