问题

给定一个巨大的概率数组和所取的百分比

probabilities = [0.1, 0.4, 0.7, 0.2, 0.9, 0.5, 0.6]
N_percentages = [20, 30, 40]  # Percentage of the list size

我想高效地计算

{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))]
        
        # Utilize previously calculated marked lists for smaller N values
        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

方案2 -使用heapq

将（idx，probability_value）Map到一个heapq，顺序为probability_value

def indicies_n_largest(values_with_indicies, percentage) -> set[int]:  # O(1) exists(int)
    """
    Returns a list of indicies for n largest probabilities in the array.

    :param arr: array of probabilities
    :param percentage: percentage of the largest probabilities to be returned
    returns: list of indicies of the largest probabilities
    """
    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, set[int]]:
    """
    Given action probabilities and a list of percentages, return a map of actions' indicies that are considered good,
    for each percentage.
    """
    values_wth_indicies = [(i, x) for i, x in enumerate(action_probs)]

    percentage_indicies_map: dict[
        int, set[int]
    ] = {}  # list of indicies of the largest probabilities

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

    return percentage_indicies_map