如何对DataFrame中按iteration分组的矩阵进行累加求和？

问题背景 #

我把多个矩阵存储在一个大型DataFrame里，示例数据如下：

import pandas as pd

data = pd.DataFrame([[13, 1, 3, 4, 0, 0], [0, 2, 6, 2, 0, 0], [3, 1, 5, 2, 2, 0], [0, 0, 10, 11, 6, 0], [5, 5, 21, 25, 41, 0], 
                     [11, 1, 3, 2, 0, 1], [3, 1, 7, 3, 1, 1], [1, 1, 6, 5, 3, 1], [1, 1, 6, 7, 6, 1], [6, 6, 21, 24, 42, 1], 
                     [17, 1, 7, 0, 0, 2], [1, 1, 6, 1, 1, 2], [2, 4, 6, 2, 1, 2], [0, 2, 11, 7, 8, 2], [5, 6, 17, 16, 46, 2], 
                     [11, 1, 10, 2, 1, 3], [2, 2, 7, 1, 1, 3], [0, 0, 14, 4, 1, 3], [0, 0, 7, 7, 5, 3], [5, 1, 20, 18, 48, 3], 
                     [16, 3, 7, 1, 2, 4], [1, 2, 4, 1, 0, 4], [2, 4, 7, 5, 3, 4], [3, 0, 4, 4, 7, 4], [7, 2, 13, 12, 58, 4]], 
                    columns=['1', '2', '3', '4', '5', 'iteration'])
print(data)

每个iteration值对应一个独立的5×5矩阵（比如iteration=0的前5行就是第一个矩阵），现在我想把这5个矩阵像普通矩阵加法那样全部加起来，得到一个最终的5×5矩阵。试了下面的代码但跑不通：

matrix = data[['1','2','3','4','5']]
print(np.sum([matrix[matrix_list['iteration']==i] for i in range(0,9)], axis=0))

问题出在哪？ #

你的代码有几个小问题：

1. 变量名写错了：matrix_list这个变量你没定义过，应该用data才对
2. 循环范围不对：你的iteration只有0到4，但你写了range(0,9)，会导致找不到对应的数据
3. 直接对DataFrame列表求和容易出问题：因为每个DataFrame的索引还是原来的行号，numpy求和时可能会因为索引不匹配导致结果混乱

几种正确的实现方法 #

方法1：用pandas分组求和（最简洁） #

因为每个矩阵的行数都是5，我们可以给每个矩阵内部的行加一个专属索引（0到4），然后按这个索引分组，把同一位置的元素跨所有iteration加起来：

import pandas as pd
import numpy as np

# 给每个矩阵内的行加索引（每个iteration内的行从0开始计数）
data['row_idx'] = data.groupby('iteration').cumcount()

# 按行索引分组，对每组的数值列求和，最后转成numpy数组就是最终矩阵
result_matrix = data.groupby('row_idx')[['1','2','3','4','5']].sum().values

print(result_matrix)

方法2：循环累加矩阵（最直观） #

如果你习惯一步步来，可以先取出第一个矩阵当初始值，然后依次把剩下的矩阵加进去：

import pandas as pd
import numpy as np

# 定义要处理的数值列
num_cols = ['1','2','3','4','5']
# 获取所有存在的iteration值
all_iterations = data['iteration'].unique()

# 初始化结果为第一个iteration对应的矩阵
sum_matrix = data[data['iteration'] == all_iterations[0]][num_cols].values

# 循环累加剩下的每个矩阵
for iter_num in all_iterations[1:]:
    current_matrix = data[data['iteration'] == iter_num][num_cols].values
    sum_matrix += current_matrix

print(sum_matrix)

方法3：修正你原来的思路 #

把你的代码改一改，解决变量名和范围的问题，同时把DataFrame转成numpy数组再求和：

import pandas as pd
import numpy as np

num_cols = ['1','2','3','4','5']
# 生成所有iteration对应的矩阵数组列表（注意iteration是0-4，所以range(5)）
matrices = [data[data['iteration'] == i][num_cols].values for i in range(5)]
# 对列表里的所有矩阵按元素位置求和
sum_matrix = np.sum(matrices, axis=0)

print(sum_matrix)

这三种方法都能得到正确的结果，你可以根据自己的习惯选一种用~

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