pandas 分组并连接列表

Question

我有一个数据表 df，里面有两列。我想根据其中一列进行分组，并把同一组的列表合并在一起，举个例子：

column_a, column_b
1,         [1,2,3]
1,         [2,5]
2,         [5,6]

处理后：

column_a, column_b
1,         [1,2,3,2,5]
2,         [5,6]

我想保留所有的重复项。我有以下几个问题：

• 这个数据表的类型都是对象（object）。使用 convert_objects() 并不能自动把 column_b 转换成列表。我该怎么做呢？
• 在 df.groupby(...).apply(lambda x: ...) 这个函数中，x 是什么？它的形式是列表吗？
• 解决我主要问题的方法是什么？

谢谢大家！

Answer 1

使用numpy和简单的“for”循环或者“map”函数：

import numpy as np

u_clm = np.unique(df.column_a.values)
all_lists = []

for clm in u_clm:
    df_process = df.query('column_a == @clm')
    list_ = np.concatenate(df.column_b.values)
    all_lists.append((clm, list_.tolist()))

df_sum_lists = pd.DataFrame(all_lists)

对于大数据集，这种方法比简单的“分组-聚合-求和”快了350倍。

Answer 2

上面提到的方法，使用 df.groupby('column_a').agg(sum) 确实可以实现。不过，你得确保你的列表里只包含 整数，否则输出的结果就会不一样。

如果你想把列表里的所有项目都转换成整数，可以使用：

df['column_a'] = df['column_a'].apply(lambda x: list(map(int, x)))

Answer 3

这个被接受的答案建议使用groupby.sum，在处理少量列表时效果不错，但用这个方法来连接多个列表的效率是平方级别的。

如果列表数量比较多，使用itertools.chain或者列表推导式会是一个更快的选择：

df = pd.DataFrame({'column_a': ['1', '1', '2'],
                   'column_b': [['1', '2', '3'], ['2', '5'], ['5', '6']]})

itertools.chain：

from itertools import chain
out = (df.groupby('column_a', as_index=False)['column_b']
         .agg(lambda x: list(chain.from_iterable(x)))
       )

列表推导式：

out = (df.groupby('column_a', as_index=False, sort=False)['column_b']
         .agg(lambda x: [e for l in x for e in l])
      )

输出结果：

  column_a         column_b
0        1  [1, 2, 3, 2, 5]
1        2           [5, 6]

速度比较

使用n次重复的例子来展示合并列表数量对速度的影响：

test_df = pd.concat([df]*n, ignore_index=True)

注意：还比较了numpy的方法（agg(lambda x: np.concatenate(x.to_numpy()).tolist())）。

Answer 4

df.groupby('column_a').agg(sum)

这个之所以有效，是因为运算符重载。sum这个函数把列表连接在了一起。最终得到的df（数据框）的索引将会是来自column_a的值：

Answer 5

object 数据类型是一种通用的数据类型，简单来说就是不属于整数、浮点数、布尔值、日期时间或时间差的类型。所以它实际上是把这些数据当作一个列表来存储。convert_objects 这个功能会尝试把某一列的数据转换成这些特定的数据类型。

你想要的是

In [63]: df
Out[63]: 
   a          b    c
0  1  [1, 2, 3]  foo
1  1     [2, 5]  bar
2  2     [5, 6]  baz


In [64]: df.groupby('a').agg({'b': 'sum', 'c': lambda x: ' '.join(x)})
Out[64]: 
         c                b
a                          
1  foo bar  [1, 2, 3, 2, 5]
2      baz           [5, 6]

这段代码是根据列 a 中的值来对数据框进行分组。想了解更多，可以查看 groupby 的相关内容。

这段代码是在进行普通的列表 sum 操作（连接），就像 [1, 2, 3] + [2, 5] 这样，结果会是 [1, 2, 3, 2, 5]。