将仅含0、0.5、1、1.5的列表高效转换为PMF的最优方法

我有多个仅包含0、0.5、1、1.5的列表，希望将每个列表高效转换为概率质量函数（PMF）。例如列表[0.5, 0.5, 1, 1.5]对应的PMF为[0, 0.5, 0.25, 0.25]。由于需要多次处理这类大型列表，若能避免循环则最佳，请问实现这一需求的最高效方法是什么？

当前低效实现 #

以下是我当前的实现方案，我认为这种方式效率低下且不够简洁：

def get_distribution(samplemodes1):
    
    n, bin_edges = np.histogram(samplemodes1, bins = 9)
    totalcount = np.sum(n)
    bin_probability = n / totalcount
    bins_per_point = np.fmin(np.digitize(samplemodes1, bin_edges), len(bin_edges)-1)
    probability_perpoint = [bin_probability[bins_per_point[i]-1] for i in range(len(samplemodes1))] 
    
    counts = Counter(samplemodes1)
    total = sum(counts.values())
    
    probability_mass = {k:v/total for k,v in counts.items()}
    #print(probability_mass)
    
    key_values = {}
    
    if(0 in probability_mass):
        key_values[0] = probability_mass.get(0)
    else:
        key_values[0] = 0
    if(0.5 in probability_mass):
        key_values[0.5] = probability_mass.get(0.5)
    else:
        key_values[0.5] = 0
    if(1 in probability_mass):
        key_values[1] = probability_mass.get(1)
    else:
        key_values[1] = 0
    if(1.5 in probability_mass):
        key_values[1.5] = probability_mass.get(1.5)  
    else:
        key_values[1.5] = 0
        
        
    distribution = list(key_values.values())
    return distribution

高效实现方案 #

方法一：Numpy bincount映射法（性能最优） #

将固定浮点值映射为整数索引，利用np.bincount一次性统计所有频次，全程向量化操作，无循环，处理大型数组速度最快：

import numpy as np

def get_pmf(arr):
    arr_np = np.array(arr)
    # 将0→0、0.5→1、1→2、1.5→3映射为整数索引
    indices = (arr_np * 2).astype(np.int32)
    # 统计每个索引的出现次数，minlength确保覆盖4个取值
    counts = np.bincount(indices, minlength=4)
    # 计算概率并转为列表
    return (counts / len(arr_np)).tolist()

方法二：简化Counter实现（标准库方案） #

利用Counter.get()的默认值特性，直接按固定顺序生成PMF，删除冗余的条件判断：

from collections import Counter

def get_pmf(arr):
    count_dict = Counter(arr)
    total = len(arr)
    # 按固定取值顺序生成概率列表，缺失值默认取0
    return [count_dict.get(val, 0) / total for val in [0, 0.5, 1, 1.5]]

方法三：Numpy count_nonzero法 #

直接统计每个固定值的出现次数，逻辑直观，同样无循环：

import numpy as np

def get_pmf(arr):
    arr_np = np.array(arr)
    total = len(arr_np)
    return [
        np.count_nonzero(arr_np == 0) / total,
        np.count_nonzero(arr_np == 0.5) / total,
        np.count_nonzero(arr_np == 1) / total,
        np.count_nonzero(arr_np == 1.5) / total
    ]

性能说明 #

• 处理百万级以上的大型列表时，Numpy bincount方法效率最高，比原始实现快数十倍；
• 原始实现中冗余的直方图计算代码完全可以删除，仅保留Counter相关逻辑也能大幅提升效率；
• 所有高效方案均避免了显式循环，充分利用Python/Numpy的内部优化机制。

内容的提问来源于stack exchange，提问作者Anthony Petruzzio