如何无循环查找二维数组各列一维数组元素索引并优化性能

核心思路 #

放弃逐个对field元素调用argwhere的列表推导，改用numpy的排序+索引映射或二分查找实现批量匹配，把O(n*m)的时间复杂度降到O(n log n)（n为参赛者规模），完全利用numpy底层的C级运算效率。

具体实现方案 #

假设你的ranks是形状为(模拟次数, 参赛者数)的二维数组，field是每次迭代中需要匹配的目标排名数组。

方案1：排序+索引映射（最优解） #

通过预排序建立排名到原始索引的映射，直接批量获取所有匹配位置：

import numpy as np

# 示例数据：10000次模拟，25000名参赛者
simulations = 10000
participants = 25000
ranks = np.random.randint(1, participants+1, size=(simulations, participants))
field = np.random.randint(1, participants+1, size=(simulations, participants))

# 对每一次模拟的ranks排序，同时记录原始索引
sorted_ranks, sorted_indices = np.sort(ranks, axis=1), np.argsort(ranks, axis=1)

# 用二分查找快速定位field元素在sorted_ranks中的位置
pos = np.searchsorted(sorted_ranks, field, side='left', axis=1)
# 过滤无效匹配（如果field存在ranks中没有的值）
valid_mask = pos < participants
pos[~valid_mask] = 0

# 映射回原始索引
result_indices = sorted_indices[np.arange(simulations)[:, None], pos]

如果你的需求是直接按排名分配积分，甚至可以跳过索引查找步骤，直接基于排序索引批量赋值：

# 示例积分规则：排名第k名得(participants - k + 1)分
scores = np.zeros_like(ranks)
# 按排序后的索引从高到低赋值积分
scores[np.arange(simulations)[:, None], sorted_indices] = np.arange(participants, 0, -1)

方案2：哈希映射（适用于有大量重复排名的场景） #

如果ranks中存在大量重复排名，可通过np.unique建立排名到索引的映射：

def get_rank_indices(rank_row):
    vals, inv = np.unique(rank_row, return_inverse=True)
    # 批量生成每个排名对应的所有原始索引
    idx_groups = np.split(np.argsort(inv), np.cumsum(np.bincount(inv))[:-1])
    return dict(zip(vals, idx_groups))

# 对每一次模拟生成映射（仍有行级循环，但重复值多时代价远低于argwhere）
rank_maps = [get_rank_indices(row) for row in ranks]
# 批量获取field对应的索引
result_indices = np.array([[rank_maps[i][f] for f in field[i]] for i in range(simulations)])

性能提升原因 #

• 排序、二分查找都是numpy底层优化的向量化操作，比Python循环快10~100倍
• 排序+查找的整体复杂度为O(n log n)，远低于列表推导+argwhere的O(n*m)
• 全程避免Python层面的循环，完全利用numpy的C语言运算效率

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