如何将嵌套循环改写为地道numpy代码生成赛事胜负频率矩阵

你要的地道Numpy矢量化实现如下，完全去掉了Python层面的循环，大输入场景下运行效率远高于原始嵌套循环实现：

import numpy as np

input_example = np.array([
    [55.90, 81.50, 76.60, 69.50],
    [52.50, 74.60, 74.00, 64.80],
    [52.40, 74.90, 78.20, 60.90],
    [52.60, 78.90, 77.60, 60.80],
])

output_example = np.array([
    [0.50, 1.00, 0.75, 0.75],
    [0.00, 0.50, 0.50, 0.25],
    [0.25, 0.50, 0.50, 0.50],
    [0.25, 0.75, 0.50, 0.50],
])

def results_matrix(tournament_results):
    players, tournaments = tournament_results.shape
    # 利用广播机制批量比较所有选手的对阵成绩
    win_count = (tournament_results[:, np.newaxis] > tournament_results).sum(axis=-1)
    win_rate = win_count / tournaments
    # 对角线（自己对阵自己）统一设为0.5
    np.fill_diagonal(win_rate, 0.5)
    return win_rate

if __name__ == '__main__':
    res_matrix = results_matrix(input_example)
    assert np.allclose(res_matrix, output_example)

实现说明 #

• 核心逻辑利用了Numpy的广播特性：我们把形状为(选手数, 赛事数)的输入数组，通过[:, np.newaxis]扩展为(选手数, 1, 赛事数)的三维数组，和原数组直接比较时会自动广播为(选手数, 选手数, 赛事数)的布尔矩阵，矩阵中[i,j,k]位置为True代表选手i在第k场赛事中成绩高于选手j
• 对最后一个维度（赛事维度）求和即可得到每个选手对阵其他选手的总胜场数，除以总赛事数得到原始胜率
• 最后用np.fill_diagonal将对角线（自己对阵自己的胜率）统一设置为0.5
• 额外修正了原实现的隐藏问题：原代码初始化结果矩阵时形状设为了(players, tournaments)，只有当选手数等于赛事数时才能正常运行，优化后的实现可适配任意选手数、赛事数的输入

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