如何定义按指定分组大小计算NumPy数组均值差的函数

直接给出完整可运行的numpy实现代码：

import numpy as np

def diff_avg(arr, size):
    unit_length = 2 * size
    # 截断不足1个计算单元的尾部元素
    valid_length = (len(arr) // unit_length) * unit_length
    if valid_length == 0:
        return []
    # 将有效元素重排为 N 个计算单元的二维结构
    valid_arr = arr[:valid_length].reshape(-1, unit_length)
    # 批量计算所有单元的均值差
    return (np.mean(valid_arr[:, :size], axis=1) - np.mean(valid_arr[:, size:], axis=1)).tolist()

测试验证 #

用你给出的示例数组测试，输出完全符合预期：

arr = np.array([1,2,3,4,5,6,7,8])
print(diff_avg(arr, size=2)) # 输出 [-2, -2]
print(diff_avg(arr, size=3)) # 输出 [-3]

关键逻辑说明 #

• 先按2*size的长度对原数组做截断，只保留刚好能组成整数个计算单元的元素
• 通过reshape(-1, unit_length)将一维数组转为二维数组，每一行就是一个独立的计算单元，无需手动循环截取后续分组
• 利用numpy的批量计算能力，直接对所有行的前size列和后size列分别求均值再做差，一次计算得到所有结果

可选：纯Python实现（无需numpy） #

如果你不想依赖numpy，也可以用原生循环实现：

def diff_avg(arr, size):
    unit_length = 2 * size
    res = []
    # 步长为2*size遍历数组
    for i in range(0, len(arr), unit_length):
        unit = arr[i:i+unit_length]
        # 不足一个单元直接终止循环
        if len(unit) < unit_length:
            break
        avg_first = sum(unit[:size]) / size
        avg_second = sum(unit[size:]) / size
        res.append(avg_first - avg_second)
    return res

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