按比例选高概率元素生成原索引0/1数组的高效实现

给定一个大规模概率数组，以及一组需要选取的比例值列表：

probabilities = [0.1, 0.4, 0.7, 0.2, 0.9, 0.5, 0.6]
N_percentages = [20, 30, 40]  # 列表规模的百分比

需要高效计算得到如下结果：

{20:[0, 0, 0, 0, 1, 0, 0], 30:[0, 0, 1, 0, 1, 0, 0], 40:[0, 0, 1, 0, 1, 0, 1]}

要求不能丢失原始索引——标记值必须保留在原数组对应的位置上。

以下是我的尝试方案：

def mark_probabilities_for_multiple_N1(probabilities, N_percentages):
    marked_lists = {}
    sorted_indices = sorted(range(len(probabilities)), key=lambda i: probabilities[i], reverse=True)
    list_size = len(probabilities)
    
    for N_percentage in N_percentages:
        N = int(N_percentage * list_size / 100)
        
        marked_lists[N_percentage] = [1 if i in sorted_indices[:N] else 0 for i in range(len(probabilities))]
        
        # 利用已计算的较小N值的标记列表
        for prev_N_percentage in [prev_N for prev_N in marked_lists if prev_N < N_percentage]:
            marked_lists[N_percentage] = [1 if marked_lists[prev_N_percentage][i] == 1 or marked_lists[N_percentage][i] == 1 else 0 for i in range(len(probabilities))]

    return marked_lists

将（索引，概率值）映射到heapq中，按probability_value排序：

def indicies_n_largest(values_with_indicies, percentage) -> list[int]:
    """
    返回数组中前n个最大概率值的索引列表。

    :param values_with_indicies: 带索引的概率数组
    :param percentage: 需返回的最大概率值的百分比
    :returns: 最大概率值的索引列表
    """
    fraction = percentage / 100
    samples_num = int(len(values_with_indicies) * fraction)
    result = heapq.nlargest(samples_num, values_with_indicies, key=lambda x: x[1])
    return [x[0] for x in result]


def percentage_indicies_map(action_probs, percentages) -> dict[int, list[int]]:
    """
    给定动作概率和比例列表，返回各比例下被判定为优质的动作索引映射表。
    """
    values_wth_indicies = [(i, x) for i, x in enumerate(action_probs)]

    percentage_indicies_map: dict[
        int, list[int]
    ] = {}  # 存储各比例对应的最大概率值索引列表

    for percentage in percentages:
        percentage_indicies_map[percentage] = indicies_n_largest(values_wth_indicies, percentage)

    return percentage_indicies_map

内容的提问来源于Stack Exchange，提问作者Bartek Lachowicz