python - 使用 numpy 数组作为条目分组并减少 pandas 数据帧

Question

我有一个具有以下结构的 pandas.DataFrame:

>>> data 
a    b    values
1    0    [1, 2, 3, 4]
2    0    [3, 4, 5, 6]
1    1    [1, 3, 7, 9]
2    1    [2, 4, 6, 8]

('values' 的类型为 numpy.array)。我想要做的是按 'a' 列对数据进行分组，然后合并值列表。我的目标是最终实现以下目标:

>>> data 
a    values
1    [1, 2, 3, 4, 1, 3, 7, 9]
2    [3, 4, 5, 6, 2, 4, 6, 8]

请注意，值的顺序无关紧要。我如何实现这一目标？我想过类似的事情

>>> grps = data.groupby(['a'])
>>> grps['values'].agg(np.concatenate)

但这失败了，出现了 KeyError。我敢肯定有一个 pandaic 的方式来实现这一点 - 但如何呢？谢谢。

Answer 1

类似于 John Galt 的回答，您可以分组然后应用 np.hstack:

In [278]: df.groupby('a')['values'].apply(np.hstack)
Out[278]: 
a
1    [1, 2, 3, 4, 1, 3, 7, 9]
2    [3, 4, 5, 6, 2, 4, 6, 8]
Name: values, dtype: object 

要取回你的框架，你需要 pd.Series.to_frame 和 pd.reset_index:

In [311]: df.groupby('a')['values'].apply(np.hstack).to_frame().reset_index()
Out[311]: 
   a                    values
0  1  [1, 2, 3, 4, 1, 3, 7, 9]
1  2  [3, 4, 5, 6, 2, 4, 6, 8]

---

性能

df_test = pd.concat([df] * 10000) # setup

%timeit df_test.groupby('a')['values'].apply(np.hstack) # mine
1 loop, best of 3: 219 ms per loop

%timeit df_test.groupby('a')['values'].sum() # John's 
1 loop, best of 3: 4.44 s per loop

sum 对于列表来说效率很低，当 Values 是一个 np.array 时不起作用。