如何在NumPy中高效获取数组A对应其唯一行数组B的索引？

嘿，针对你这个处理大规模二维数组、要快速得到A每行对应其排序唯一行数组B的索引的问题，我有个完全矢量化的方案——全程不用循环，靠NumPy底层优化的操作搞定，完美适配大规模数据，还兼容NumPy 1.1版本。

核心思路 #

二维数组的行没法直接用一维数组的索引查找方法，那我们就把每一行打包成一个不可分割的np.void类型标量，把二维问题转化为一维问题，这样就能用排序去重和np.searchsorted这些高效的矢量化工具了。

具体实现步骤 #

1. 构造示例数组（模拟你的场景）

import numpy as np

# 示例二维数组A
A = np.array([[1, 2], [3, 4], [1, 2], [5, 6], [3, 4]])
# 预期的B是排序后的唯一行：[[1,2], [3,4], [5,6]]

2. 将A的行打包成np.void一维数组

这一步是关键，把每行转换成一个单一的“原子”元素：

# 计算每行对应的字节长度（根据数组的 dtype 和列数）
row_byte_size = A.dtype.itemsize * A.shape[1]
# 将A的行视图转换为一维void数组
A_void = A.view(f'V{row_byte_size}').ravel()

3. 生成排序后的唯一行数组B（兼容NumPy 1.1）

因为NumPy 1.1的np.unique还没有return_index参数，我们用排序+掩码的方式生成B：

# 先对void数组和原数组按void元素排序
sorted_idx = np.argsort(A_void)
sorted_A_void = A_void[sorted_idx]
sorted_A = A[sorted_idx]

# 生成唯一元素的掩码：只保留和前一个元素不同的位置
unique_mask = np.concatenate([[True], sorted_A_void[1:] != sorted_A_void[:-1]])

# 得到排序后的唯一行数组B，以及对应的void数组
B = sorted_A[unique_mask]
B_void = sorted_A_void[unique_mask]

4. 快速查找A每行对应的B的索引

用np.searchsorted（NumPy 1.1支持这个函数）一次性得到所有索引：

# 查找每个A_void元素在B_void中的位置，就是A每行对应B的索引
indices = np.searchsorted(B_void, A_void)

验证结果

print("原数组A:\n", A)
print("排序后的唯一行数组B:\n", B)
print("A每行对应的B的索引:", indices)

输出结果：

原数组A:
 [[1 2]
 [3 4]
 [1 2]
 [5 6]
 [3 4]]
排序后的唯一行数组B:
 [[1 2]
 [3 4]
 [5 6]]
A每行对应的B的索引: [0 1 0 2 1]

为什么这个方案高效？ #

所有操作都是NumPy底层矢量化实现，没有Python层面的循环，处理百万级甚至更大的数组时，速度比循环快几个数量级——完全符合你处理超大规模数组的需求。

注意事项 #

• 这个方法要求数组A是同质的（所有元素 dtype 一致），这也是NumPy数组的默认特性；
• np.void的转换基于字节，所以A的每行长度必须相同（二维数组天然满足）。

内容的提问来源于stack exchange，提问作者Koyagi