为何第二种numpy实现决策树熵计算方法更快？能否优化第一种？

我在开发决策树时需要计算numpy数组标签向量y的熵，试过两种实现方式：

第一种（生成器表达式求和）：

entropy = sum(-len(y[y == label]) / len(y) * np.log2(len(y[y == label]) / len(y)) for label in labels)

第二种（循环累加）：

entropy = 0
for label in labels:
    cond = y == label 
    fraction = len(y[cond]) / len(y)
    entropy += -fraction * np.log2(fraction)
return entropy

注：labels是np.unique(y)的结果

实际计时后发现第二种方法更快，我猜测是第一种重复计算了分数值，具体原因是什么？能不能优化第一种方法提升速度？

性能差异的核心原因 #

第一种方法的问题确实是重复执行了相同的数组运算：

• 对每个label，len(y[y == label]) / len(y)这个分数被计算了两次——一次用于乘法项，一次用于np.log2()的参数。
• 每次计算都要执行y == label的布尔数组生成，以及y[cond]的索引和长度统计，这两步是numpy里相对耗时的操作，重复执行直接翻倍了这部分的开销。
• 第二种方法把分数存入变量fraction，只计算一次布尔索引和长度，后续直接复用变量值，避免了冗余计算，因此速度更快。

优化第一种方法的方案 #

可以通过避免重复计算分数来优化第一种实现，有两种常见方式：

方式1：使用海象运算符（Python 3.8+）

在生成器表达式里先计算一次分数，复用这个值：

entropy = sum(-frac * np.log2(frac) for label in labels if (frac := len(y[y == label]) / len(y)) > 0)

这里用海象运算符:=提前计算frac，确保每个label只执行一次y == label和长度统计，同时加入frac > 0的判断避免np.log2(0)的报错（虽然labels是np.unique(y)的结果，理论上不会出现0，但增加判断更健壮）。

方式2：先预计算所有分数再求和

把分数的计算和熵的计算分开，更直观：

total = len(y)
fractions = [len(y[y == label]) / total for label in labels]
entropy = sum(-f * np.log2(f) for f in fractions)

这种方式和第二种循环的逻辑本质一致，只是用列表推导式预生成所有分数，再用生成器求和，既保留了第一种的简洁性，又消除了重复计算。

内容的提问来源于stack exchange，提问作者Python Snake