python - 如何加速 Pandas 操作(按嵌套列表项分组)

Question

假设我有以下df:

   col1  col2    col3
0     1     3    brazil argentina spain uk
1     2     4    brazil france iceland
2     7     9    iceland brazil
3     8     1    chile
4     5     0    spain italy

我可以通过创建一个集合来获取 col3 中唯一单词主题的列表:

split_list = []
for i in df['col3'].to_list()
    themes = i.split(' ')
    for j in themes:
        split_list.append(j)
unique_themes = list(set(split_list))

其中提供以下内容:

unique_themes = ['brazil', 'argentina', 'spain', 'uk', 'france', 'iceland', 'chile', 'italy']

然后我可以创建新的数据框，为 unique_themes 中的每个主题提供 col1 和 col2 的总和:

theme_df = pd.DataFrame()

columns = df.columns.tolist()
columns.remove('col3')

for theme in unique_themes:

    data = df[df['col3'].apply(lambda x: theme in x.split(' '))][columns]
    data["theme"] = theme
    data["count"] = 1
    theme_df = pd.concat([theme_df, data], axis=0)

data = theme_df.groupby("theme", as_index=False)[["theme"] + ["count"] + columns].sum()

给出以下theme_df:

      theme      count   col1   col2
0     brazil         3     10     16 
1     argentina      1      1      3
2     spain          2      6      3
3     uk             1      1      3
4     france         1      2      4
5     iceland        2      9     13
6     chile          1      8      1
7     italy          1      5      0

但是这非常慢。在包含 15,000 行、大约 5,000 个独特单词主题的数据帧上，for 循环(检查 theme 是否包含在 col3 中) )大约需要 2 分钟来计算。

有更好(更快)的方法吗？也许首先对数据进行规范化，以便 col3 仅包含一个主题，从而无需 apply？或者完全是另一种方法？

Answer 1

使用DataFrame.explode对于由 Series.str.split 转换为列表的重复值然后按名称聚合 GroupBy.agg 进行聚合:

df1 = (df.assign(col3 = df['col3'].str.split())
         .explode('col3')
         .groupby('col3', sort=False)
         .agg(count=('col1','size'),
              col1=('col1','sum'),
              col2=('col2','sum'))
         .reset_index())
print (df1)
        col3  count  col1  col2
0     brazil      3    10    16
1  argentina      1     1     3
2      spain      2     6     3
3         uk      1     1     3
4     france      1     2     4
5    iceland      2     9    13
6      chile      1     8     1
7      italy      1     5     0

编辑:您可以动态创建元组字典，添加计数的第一个键并将 ** 传递给 agg:

cols = df.columns.difference(['col3'])
d = {c: (c,'sum') for c in cols}
#merge dictionaries
d = {**{'count': (cols[0],'size')}, **d}
print (d)
{'count': ('col1', 'size'), 'col1': ('col1', 'sum'), 'col2': ('col2', 'sum')}

df1 = (df.assign(col3 = df['col3'].str.split())
        .explode('col3')
        .groupby('col3', sort=False)
        .agg(**d)
        .reset_index())
print (df1)
        col3  count  col1  col2
0     brazil      3    10    16
1  argentina      1     1     3
2      spain      2     6     3
3         uk      1     1     3
4     france      1     2     4
5    iceland      2     9    13
6      chile      1     8     1
7      italy      1     5     0