Python实现矩阵指定范围元素几何均值的两种方案求解

实现代码 #

直接用numpy的布尔索引替代嵌套循环，性能远高于循环实现，同时和原有逻辑完全对齐：

import numpy as np

def mygmean_vectorized(array, t):
    """
    计算矩阵中指定范围内元素的几何均值，返回几何均值与范围外元素组成的数组
    """
    # 构造布尔掩码，筛选处于[t[0], t[-1]]区间的元素
    in_range = (array >= t[0]) & (array <= t[-1])
    # 分别提取范围内、范围外元素
    included = array[in_range]
    excluded = array[~in_range]
    # 计算几何均值，逻辑和循环版本完全一致
    geomean = np.prod(included) ** (1/len(included))
    return geomean, excluded

测试验证 #

运行测试用例：

print(mygmean_vectorized(np.array([[2,4],[-9,3],[6,-2],[8,1],[12,8]]),(0,10)))

输出结果和预期完全一致：

(3.684369762785971, array([-9, -2, 12]))

补充优化建议 #

如果处理的数组数值很大、元素很多，直接计算乘积可能出现数值溢出，可改用对数法计算更稳定的几何均值，结果和原有计算逻辑一致：

geomean = np.exp(np.mean(np.log(included)))

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