首页 > 解决方案 > groupby 当列包含需要添加的多维数组时

问题描述

我有一个包含二维向量作为列的熊猫数据框。我想按其中一列分组并将向量相加。

我已经尝试 groupby 然后 sum ,如下面的代码所示,但输出列是向向量添加维度而不是添加向量(类似于使用 np.add 时)。

import pandas as pd

data = pd.DataFrame({'label': ['A', 'B', 'A'], 'label2' : ['X', 'Y', 'Z'],
                     'output' :   [[[1,2,3,4],[5,6,7,8]] ,[[9,10,11,12],[13,14,15,16]],[[17,18,19,20],[21,22,23,24]]] })

data_grouped = data.groupby('label')['output'].sum()

我想按“标签”分组并汇总输出。鉴于输出是二维向量,我希望向量被添加而不是组合。因此,我的期望是:

label A: output is [[18,20,22,24],[26,28,30,32]]
label B: output is [[9,10,11,12],[13,14,15,16]]

但我得到:

label A: [[1, 2, 3, 4], [5, 6, 7, 8], [17, 18, 19, 20],[21,22,23,24]]

label B: [[9, 10, 11, 12], [13, 14, 15, 16]]

标签: python-3.xpandasnumpy

解决方案


解决方案

import pandas as pd
import numpy as np

data = pd.DataFrame({'label': ['A', 'B', 'A'], 'label2' : ['X', 'Y', 'Z'],
                     'output' :   [[[1,2,3,4],[5,6,7,8]] ,[[9,10,11,12],[13,14,15,16]],[[17,18,19,20],[21,22,23,24]]] })

data['output'] = data['output'].map(np.array)
data_grouped = data[['label', 'output']].groupby('label').sum()
print(data_group)

>>>                                      output
>>> label                                      
>>> A      [[18, 20, 22, 24], [26, 28, 30, 32]]
>>> B       [[9, 10, 11, 12], [13, 14, 15, 16]]

说明

output包含 python 列表。对 2 个列表的操作+将列表连接在一起:

print([1, 2] + [3, 4])
>>> [1, 2, 3, 4]

print([[1], [2]] + [[3], [4]])
>>> [[1], [2], [3], [4]]

data['output'].map(np.array)将您的 2D 列表转换为 2D numpy 数组。Numpy 数组+操作(由 使用sum())对两个数组中“相同位置”的值求和。


推荐阅读