用 Python + NumPy 编写不带 for 循环的分组排名

Question

我的目标是创建一个函数，该函数接受两个输入数组并输出组内的排名。更具体地说，函数的第一个输入是分数数组 (

scores

，第二个输入数组是组大小数组 (

groups

，这样组大小的总和等于总金额)的分数，并且分数数组可以分为类似

np.split(scores, np.cumsum(groups))

的组。对于组内得分最高的点，输出排名应该为 1，对于第二高的点，输出排名应该为 2，依此类推。我已经实现了这个如下（包括两个测试用例）：

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()

虽然这可以正常工作，但我想知道是否可以在不使用 for 循环的情况下实现这一点，从而可能提高性能。我进行了多次尝试但没有成功。有什么建议吗？

我尝试过其他方法，例如利用

np.lexsort

和

np.split

。没有任何效果。

Answer 1

如果所有值都是正数，则可以将它们偏移一个与它们所在组索引的倍数相对应的基值。然后使用 argsort 获取具有偏移量的项目的全局位置。这将使给定组的成员相对于组索引保持在一起。然后将全局索引转换回组相对位置并反转顺序以获得排名：

def assign_ranks(scores,groups):
    base      = np.repeat(np.arange(groups.size),groups)
    order     = np.argsort(scores+base*np.max(scores))
    groupBase = np.repeat(np.cumsum(np.insert(groups[:-1],0,0)),groups)
    return np.repeat(groups,groups) - order + groupBase