如何在Pandas DataFrame中按reference和subject列分组聚合

问题描述 #

现有如下结构的DataFrame：

page    reference       ids                 subject           word
1       apple           ['aaaa', 'bbbbb', 'cccc']       name            app
1       apple           ['bndv', 'asasa', 'swdsd']      fruit           is
1       apple           ['bsnm', 'dfsd', 'dgdf']        fruit           text
1       bat             ['asas', 'ddfgd', 'ff']         thing           sport
1       cat             ['sds', 'dffd', 'gdg']          fruit           color
1       bat             ['sds', 'fsss', 'ssfd']         thing           was
1       bat             ['fsf', 'sff', 'fss']           place           that
2       dog             ['fffds', 'gd', 'sdg']          name            mud
2       egg             ['dfff', 'sdf', 'vcv']          place           gun
2       dog             ['dsfd', 'fds', 'gfdg']         thing           kit
2       egg             ['ddd', 'fg', 'dfg']            place           hut

需要按reference和subject列分组，将每组的ids和word按指定格式聚合，最终得到如下结果：

page    reference   ids                                                subject          word
1       apple   [['bndv', 'asasa', 'swdsd'],['bsnm', 'dfsd', 'dgdf']]   fruit           [[is], [text]]
1       apple   ['aaaa', 'bbbbb', 'cccc']                               name            [app]
1       bat     [['asas', 'ddfgd', 'ff'], ['sds', 'fsss', 'ssfd']]     thing           [[sport], [was]]
1       bat     ['fsf', 'sff', 'fss']                                   place           [that]
1       cat     ['sds', 'dffd', 'gdg']                                  fruit           [color]
2       dog     ['fffds', 'gd', 'sdg']                                  name            [mud]
2       dog     ['dsfd', 'fds', 'gfdg']                                 thing           [kit]
2       egg     [['dfff', 'sdf', 'vcv'], ['ddd', 'fg', 'dfg']]          place           [[gun], [hut]]

代码实现 #

import pandas as pd

# 构造原始DataFrame
data = {
    'page': [1,1,1,1,1,1,1,2,2,2,2],
    'reference': ['apple','apple','apple','bat','cat','bat','bat','dog','egg','dog','egg'],
    'ids': [
        ['aaaa', 'bbbbb', 'cccc'],
        ['bndv', 'asasa', 'swdsd'],
        ['bsnm', 'dfsd', 'dgdf'],
        ['asas', 'ddfgd', 'ff'],
        ['sds', 'dffd', 'gdg'],
        ['sds', 'fsss', 'ssfd'],
        ['fsf', 'sff', 'fss'],
        ['fffds', 'gd', 'sdg'],
        ['dfff', 'sdf', 'vcv'],
        ['dsfd', 'fds', 'gfdg'],
        ['ddd', 'fg', 'dfg']
    ],
    'subject': ['name','fruit','fruit','thing','fruit','thing','place','name','place','thing','place'],
    'word': ['app','is','text','sport','color','was','that','mud','gun','kit','hut']
}
df = pd.DataFrame(data)

# 定义聚合函数
def agg_ids(x):
    # 组内多个元素时嵌套成列表，单个元素直接返回原列表
    return list(x) if len(x) > 1 else x.iloc[0]

def agg_word(x):
    # 组内多个元素时每个单词单独嵌套列表，单个元素直接包裹成列表
    return [[item] for item in x] if len(x) > 1 else [x.iloc[0]]

# 分组聚合并调整列顺序
result = df.groupby(['reference', 'subject'], as_index=False).agg(
    page=('page', 'first'),
    ids=('ids', agg_ids),
    word=('word', agg_word)
)[['page', 'reference', 'ids', 'subject', 'word']]

# 打印结果
print(result)

说明 #

• 分组逻辑：以reference和subject为分组键，确保同组数据是同一参照项+同一主题的记录。
• 列处理：
  • page：同组内page值一致，直接取第一个值即可。
  • ids：组内有多条记录时，将每个ids列表作为子元素合并成大列表；单条记录时保留原列表格式。
  • word：组内有多条记录时，每个单词单独用列表包裹后再合并；单条记录时直接将单词放入列表。

内容的提问来源于stack exchange，提问作者Anonymous