如何加速Numpy数组反向索引的构建？

这个场景我太熟了——当n和m的取值很大时，Python层的显式遍历循环会成为性能瓶颈，完全可以用numpy的向量化操作来彻底解决！我们可以利用numpy的扁平化、排序、分组统计等原生优化操作，把整个过程从O(n)的Python循环变成底层C实现的向量化计算，速度能提升几个数量级。

优化方案代码 #

import numpy as np

# 示例输入
a = np.array([
 [0, 1, 2],
 [0, 1, 3],
 [2, 3, 4],
 [4, 5, 6],
])
m = a.max()  # 如果已知m可以直接指定，这里自动获取

# 1. 生成所有元素对应的行索引、列索引和元素值
n_rows = a.shape[0]
row_indices = np.repeat(np.arange(n_rows), 3)  # 每行重复3次，对应三个列
col_indices = np.tile(np.arange(3), n_rows)    # 列索引循环[0,1,2]
flat_a = a.flatten()                           # 扁平化数组

# 2. 按元素值排序，让相同元素聚在一起
sorted_indices = np.argsort(flat_a)
sorted_vals = flat_a[sorted_indices]
sorted_rows = row_indices[sorted_indices]
sorted_cols = col_indices[sorted_indices]

# 3. 找到每个元素的起始位置和出现次数
unique_vals, idx_start, counts = np.unique(sorted_vals, return_counts=True, return_index=True)

# 4. 初始化结果数组，默认填充-1
b_row = -np.ones((m+1, 2), dtype=int)
b_col = -np.ones((m+1, 2), dtype=int)

# 5. 填充第一次出现的位置
b_row[unique_vals, 0] = sorted_rows[idx_start]
b_col[unique_vals, 0] = sorted_cols[idx_start]

# 6. 填充第二次出现的位置（仅对出现2次的元素）
mask = counts == 2
b_row[unique_vals[mask], 1] = sorted_rows[idx_start[mask] + 1]
b_col[unique_vals[mask], 1] = sorted_cols[idx_start[mask] + 1]

# 验证结果
print("b_row:\n", b_row)
print("b_col:\n", b_col)

代码细节解释 #

• 步骤1：先把二维数组a的行、列和元素都展开成一维数组，让每个元素都对应唯一的row_indices和col_indices，方便后续统一处理。
• 步骤2：通过排序把相同元素集中在一起，这样后续可以批量处理同一元素的所有出现位置，替代原代码中逐行判断的逻辑。
• 步骤3：用np.unique一次性获取唯一元素、每个元素在排序后数组的起始索引和出现次数，这是整个向量化实现的核心，彻底避开了Python循环。
• 步骤4-6：初始化结果数组后，先填充所有元素第一次出现的位置，再对出现两次的元素填充第二次的位置，未出现的位置保持初始值-1，完全符合需求。

性能对比 #

这个方案完全避免了Python层的显式循环，所有操作都是numpy底层的C级别的计算，当n和m很大时（比如n=1e6），速度会比原代码快几十甚至上百倍。你可以用timeit模块自行测试两者的耗时差异。

内容的提问来源于stack exchange，提问作者Nico Schlömer