Pandas 中 groupby 和聚合的更快解决方案

Question

我有一个看起来像这样的数据框：

df = pd.DataFrame.from_dict({"id": [1,1,1,2,3,3], "s1": [4,np.nan,np.nan,1,1,5], "s2": [14,np.nan,19,np.nan,18,19]})

    id  s1  s2
0   1   4.0 14.0
1   1   NaN NaN
2   1   NaN 19.0
3   2   1.0 NaN
4   3   1.0 18.0
5   3   5.0 19.0

我的目标是将它转换成这样的结果

    s1          s2
id      
1   [4.0]       [14.0, 19.0]
2   [1.0]       [0]
3   [1.0, 5.0]  [18.0, 19.0]

这意味着我想按字段“id”进行分组，将所有其他列聚合到一个列表中（如果该组合的值为 NaN，则作为后备创建一个为零的列表）

我目前的纯 Pandas 方法对于较大的数据帧（> 100k 行和 ~100 列）非常慢。

def _aggregate_listwise(x):
    return list(x.dropna()) if not x.isnull().all() else [0]

df.groupby("id").agg(lambda x: _aggregate_listwise(x))

有没有一种不错的 Pythonic 方式来加快我的解决方案速度？谢谢你的建议！

Answer 1

类似于stack 然后unstack，但是对于对象列 agg，运行时间总是比平时长。

out = df.set_index('id').stack().groupby(level=[0,1]).agg(list).unstack(fill_value=[0])
Out[617]: 
            s1            s2
id                          
1        [4.0]  [14.0, 19.0]
2        [1.0]           [0]
3   [1.0, 5.0]  [18.0, 19.0]