如何优化Python互信息计算代码以解决nose.tools.nontrivial.TimeExpired超时问题？

我明白你遇到的超时问题——这段代码里的双重Python循环是最大的性能杀手，尤其是当输入数组arr规模较大时，Python的逐元素循环速度完全跟不上需求。咱们可以用NumPy的向量化操作彻底重构这段代码，把速度提升几个数量级。

原代码的问题分析 #

原代码里的两层for循环逐个遍历数组元素，还要每次判断arr[i,j] != 0，这种方式在数组元素较多时（比如几千行几千列），会花费大量时间在Python的循环逻辑上，而不是利用NumPy的底层优化能力。

优化后的代码实现 #

这里有两个高效的版本，都能彻底去掉循环：

版本1：使用掩码筛选非零元素

import numpy as np

def mutual_information(arr):
    total = arr.sum()
    # 计算行和列的边际概率，keepdims保持维度方便广播
    row_sums = arr.sum(axis=1, keepdims=True)
    col_sums = arr.sum(axis=0, keepdims=True)
    
    # 用广播计算每个元素的分母（行和*列和）
    denominator = row_sums * col_sums
    # 只处理非零元素，避免log(0)和除零错误
    non_zero_mask = arr != 0
    # 批量计算所有非零元素的贡献
    mi_sum = np.sum(arr[non_zero_mask] * np.log(arr[non_zero_mask] / denominator[non_zero_mask]))
    # 归一化后返回
    return mi_sum / total

版本2：用np.where简化逻辑（更简洁）

import numpy as np

def mutual_information(arr):
    total = arr.sum()
    row_sums = arr.sum(axis=1, keepdims=True)
    col_sums = arr.sum(axis=0, keepdims=True)
    
    # np.where自动处理零元素，将其贡献设为0
    terms = np.where(arr != 0, arr * np.log(arr / (row_sums * col_sums)), 0)
    return terms.sum() / total

优化核心点 #

• 向量化运算：用NumPy的广播机制替代Python循环，让底层的C代码处理批量计算，速度远超Python循环。
• 避免冗余判断：通过掩码或np.where一次性处理所有非零元素，不用逐元素判断。
• 维度保持：sum时用keepdims=True，让行和/列和的数组维度和原数组匹配，直接广播相乘，不用手动调整形状。

性能测试参考 #

如果用一个1000×1000的随机数组测试，原代码可能需要几十秒甚至更久，而优化后的代码只需要几毫秒就能完成计算，完全能满足2秒的时间限制。

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