如何在Python+NumPy中实现无循环的组内排名？

需求回顾 #

需要实现一个函数，接收两个输入数组：

• scores：分数数组
• groups：组大小数组，元素之和等于分数总数，可通过np.split(scores, np.cumsum(groups))完成分组

输出规则：组内最高分排名为1，次高分排名为2，依此类推。现有基于for循环的实现，希望优化为无循环版本以提升性能。

现有循环实现代码 #

import numpy as np

def assign_ranks(scores, groups):
    # Create an array to store the ranks
    ranks = np.zeros(len(scores), dtype=int)

    # Iterate through groups
    start = 0
    for group_size in groups:
        end = start + group_size
        ranks[start:end] = (-scores[start:end]).argsort().argsort() + 1
        start = end

    return ranks

def test():
    # Test case 1
    scores1 = np.array([0.1, 0.5, 0.3, 0.4, 0.5])
    groups1 = np.array([3, 2])
    result1 = assign_ranks(scores1, groups1)
    print("Test case 1 result:", result1)  # Expected: [3 1 2 2 1]
    assert np.all(result1 == np.array([3, 1, 2, 2, 1]))

    # Test case 2
    scores2 = np.array([0.7, 0.3, 0.6, 0.1, 0.2, 0.5, 0.8, 0.9, 0.4])
    groups2 = np.array([4, 3, 2])
    result2 = assign_ranks(scores2, groups2)
    print("Test case 2 result:", result2)  # Expected: [1 3 2 4 3 2 1 1 2]
    assert np.all(result2 == np.array([1, 3, 2, 4, 3, 2, 1, 1, 2]))

if __name__ == "__main__":
    test()

无循环向量化实现方案 #

实现代码 #

import numpy as np

def assign_ranks_vectorized(scores, groups):
    # 生成每个元素对应的组ID数组
    group_ids = np.repeat(np.arange(len(groups)), groups)
    # 按组ID升序、分数降序排序，得到排序后的索引序列
    sorted_indices = np.lexsort((-scores, group_ids))
    # 获取排序后的组ID序列
    sorted_group_ids = group_ids[sorted_indices]
    # 为每个组生成1到组大小的排名序列，拼接成整体排名
    group_counts = np.bincount(sorted_group_ids)
    group_ranks = np.concatenate([np.arange(1, cnt + 1) for cnt in group_counts])
    # 将排名映射回原数组的对应位置
    ranks = np.empty_like(scores, dtype=int)
    ranks[sorted_indices] = group_ranks
    return ranks

def test_vectorized():
    # Test case 1
    scores1 = np.array([0.1, 0.5, 0.3, 0.4, 0.5])
    groups1 = np.array([3, 2])
    result1 = assign_ranks_vectorized(scores1, groups1)
    print("Test case 1 vectorized result:", result1)  # Expected: [3 1 2 2 1]
    assert np.all(result1 == np.array([3, 1, 2, 2, 1]))

    # Test case 2
    scores2 = np.array([0.7, 0.3, 0.6, 0.1, 0.2, 0.5, 0.8, 0.9, 0.4])
    groups2 = np.array([4, 3, 2])
    result2 = assign_ranks_vectorized(scores2, groups2)
    print("Test case 2 vectorized result:", result2)  # Expected: [1 3 2 4 3 2 1 1 2]
    assert np.all(result2 == np.array([1, 3, 2, 4, 3, 2, 1, 1, 2]))

if __name__ == "__main__":
    test_vectorized()

核心原理说明 #

1. 组ID生成：通过np.repeat为每个元素标记所属组的编号，将分组信息转化为可向量化处理的数组。
2. 多关键字排序：使用np.lexsort先按组ID分组，同组内按分数降序排列，得到全局排序后的索引。
3. 组内排名生成：利用np.bincount统计每个组的元素数量，为每个组生成连续的排名序列（1到组大小），再拼接成整体排名。
4. 映射回原位置：通过排序后的索引，将生成的排名赋值回原数组的对应位置，得到最终结果。

性能优势 #

该方案完全基于numpy的向量化操作，避免了Python层面的循环开销，在处理大规模数据集（组数多、单组元素量大）时，性能会显著优于循环版本。

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