如何对二维NumPy数组的多行分组求和生成新数组？

你当前的嵌套循环写法效率较低，尤其是数据量较大时。下面提供两种更高效的实现方式，完全避免手动循环，适配你的数据结构需求：

方法一：NumPy原生重塑+轴求和 #

利用NumPy的数组重塑和批量求和能力，直接完成分组计算：

import numpy as np
import pandas as pd

def read_results(N, t_eval):
    # 读取CSV并转为NumPy数组
    solution_cells = pd.read_csv("battery_solution_6.csv").to_numpy()
    # 移除不需要的i列（原数据中i是第二列，索引为1）
    filtered_data = solution_cells[:, [0, 2, 3, 4]]
    # 将数组重塑为(分组数, N, 列数)的三维结构（需保证总行数是N的整数倍）
    grouped_data = filtered_data.reshape(-1, N, filtered_data.shape[1])
    # 对每组内的行求和，得到(分组数, 列数)的结果数组
    summed_results = grouped_data.sum(axis=1)
    # 按需筛选t_eval指定的时间点
    if t_eval:
        t_column = summed_results[:, 0]
        match_mask = np.isin(t_column, t_eval)
        summed_results = summed_results[match_mask]
    return summed_results

方法二：Pandas分组求和（更直观） #

既然已经用Pandas读取文件，直接按t列分组求和，代码可读性更高：

import pandas as pd

def read_results(N, t_eval):
    # 读取CSV为DataFrame，保留列名便于操作
    df = pd.read_csv("battery_solution_6.csv")
    # 按t列分组，对x、y、z列求和，最后重置索引保留t列
    summed_df = df.groupby('t')[['x', 'y', 'z']].sum().reset_index()
    # 按需筛选指定时间点
    if t_eval:
        summed_df = summed_df[summed_df['t'].isin(t_eval)]
    # 转为NumPy数组（如果需要）
    return summed_df.to_numpy()

核心优势说明 #

• 两种方法都基于底层C实现的批量运算，比Python嵌套循环快几个数量级，数据量越大优势越明显
• 自动适配整型/浮点型混合的列，求和后会自动统一为合适的数值类型
• Pandas方法更易维护，适合需要后续数据筛选、列操作的场景；NumPy方法更适合纯数值数组的高性能计算

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