基于Pandas实现文档词汇变更数据分组汇总的方法问询

我有如下Pandas DataFrame：

| word_number | change_type | doc1_paragraph_number | doc2_paragraph_number |
|-------------|-------------|-----------------------|-----------------------|
| 1           | -1          | 0                     | -1                    |
| 2           | 0           | 0                     | 0                     |
| 3           | 0           | 1                     | 1                     |
| 4           | 0           | 1                     | 1                     |
| 5           | 1           | 1                     | 1                     |
| 6           | 1           | 1                     | 1                     |
| 7           | 1           | -1                    | 1                     |
| 8           | 0           | 1                     | 1                     |
| 9           | 1           | -1                    | 1                     |
| 10          | 0           | 1                     | 1                     |
| 11          | 0           | 2                     | 2                     |
| 12          | 1           | 2                     | 2                     |
| 13          | -1          | 2                     | -1                    |
| 14          | 1           | 2                     | 2                     |
| 15          | 0           | 3                     | 3                     |
| 16          | 1           | -1                    | 3                     |
| 17          | 1           | -1                    | 3                     |
| 18          | 0           | 3                     | 3                     |
| 19          | 0           | 3                     | 3                     |
| 20          | 0           | 3                     | 3                     |
| 21          | 0           | 3                     | 3                     |
| 22          | 1           | -1                    | 3                     |
| 23          | 1           | -1                    | 3                     |
| 24          | 0           | 4                     | 4                     |
| 25          | 0           | 4                     | 4                     |
| 26          | 1           | -1                    | 5                     |
| 27          | 1           | -1                    | 5                     |
| 28          | 0           | 4                     | 5                     |
| 29          | 0           | 4                     | 5                     |
| 30          | 1           | -1                    | 5                     |
| 31          | -1          | 4                     | -1                    |
| 32          | 0           | 5                     | 6                     |
| 33          | 0           | 5                     | 6                     |
| 34          | 0           | 5                     | 6                     |
| 35          | -1          | 5                     | -1                    |
| 36          | -1          | 5                     | -1                    |
| 37          | 1           | -1                    | 7                     |
| 38          | 1           | -1                    | 7                     |
| 39          | 0           | 6                     | 7                     |
| 40          | 0           | 6                     | 7                     |

需要开发一个函数，将word_number合并为列表后分别存入doc1和doc2列，同时保留change_type列标识文档变更类型，期望得到的DataFrame如下：

| doc1                | doc2                  | change_type |
|---------------------|-----------------------|-------------|
| [1]                 | []                    | -1          |
| [2]                 | [2]                   | 0           |
| [3, 4]              | [3, 4]                | 0           |
| [5, 6, 8]           | [5, 6, 7, 8, 9]       | 2           |
| [10]                | [10]                  | 0           |
| [11]                | [11]                  | 0           |
| [12, 13, 14]        | [12, 14]              | 2           |
| [15]                | [15]                  | 0           |
| []                  | [16, 17]              | 1           |
| [18, 19, 20, 21]    | [18, 19, 20, 21]      | 0           |
| []                  | [22, 23]              | 1           |

规则说明：

1. 段落编号为-1表示该词汇未出现在对应段落中。
2. 变更类型：无变更=0，插入=1，删除=-1；word_number为每个词汇的唯一整数标识。

注：若存在连续word_number且变更类型为0（无变更），需将其合并为单个列表。例如：doc1列表包含词汇编号5、6、8，doc2列表包含5、6、7、8、9。

实现思路 #

1. 标记分组边界：通过构建分组键区分不同的变更场景，同一组的行需属于同一段落的同一变更类型（或混合变更场景）。当段落信息或变更类型的组合发生变化时，生成新的分组ID。
2. 生成文档词汇列表：对每个分组，筛选出对应文档段落编号不为-1的word_number，组成列表；段落编号为-1的词汇不加入对应文档的列表。
3. 确定分组变更类型：若分组内变更类型统一，则直接取该类型；若同时存在插入和删除（混合变更），则标记为2。

代码实现 #

import pandas as pd

def merge_word_numbers(df):
    # 构建分组键：区分同一段落的变更场景
    df['para_key'] = df.apply(lambda x: 
        f"{x['doc1_paragraph_number']}_{x['doc2_paragraph_number']}" if x['change_type'] == 0 
        else f"{max(x['doc1_paragraph_number'], x['doc2_paragraph_number'])}", axis=1)
    
    # 生成分组ID：当分组键或变更类型变化时，分组ID递增
    df['group_id'] = (df[['para_key', 'change_type']] != df[['para_key', 'change_type']].shift()).any(axis=1).cumsum()
    
    # 定义聚合逻辑
    def agg_group(group):
        # 生成doc1的词汇列表：筛选段落编号不为-1的word
        doc1_list = group[group['doc1_paragraph_number'] != -1]['word_number'].tolist()
        # 生成doc2的词汇列表：筛选段落编号不为-1的word
        doc2_list = group[group['doc2_paragraph_number'] != -1]['word_number'].tolist()
        
        # 确定分组的最终变更类型
        unique_types = group['change_type'].unique()
        if len(unique_types) == 1:
            final_type = unique_types[0]
        else:
            # 混合插入和删除，标记为2
            final_type = 2
        
        return pd.Series([doc1_list, doc2_list, final_type], index=['doc1', 'doc2', 'change_type'])
    
    # 按分组ID聚合得到结果
    result = df.groupby('group_id').apply(agg_group).reset_index(drop=True)
    return result

# 测试代码
# 构造原始DataFrame
raw_data = [
    [1, -1, 0, -1], [2, 0, 0, 0], [3, 0, 1, 1], [4, 0, 1, 1],
    [5, 1, 1, 1], [6, 1, 1, 1], [7, 1, -1, 1], [8, 0, 1, 1],
    [9, 1, -1, 1], [10, 0, 1, 1], [11, 0, 2, 2], [12, 1, 2, 2],
    [13, -1, 2, -1], [14, 1, 2, 2], [15, 0, 3, 3], [16, 1, -1, 3],
    [17, 1, -1, 3], [18, 0, 3, 3], [19, 0, 3, 3], [20, 0, 3, 3],
    [21, 0, 3, 3], [22, 1, -1, 3], [23, 1, -1, 3], [24, 0, 4, 4],
    [25, 0, 4, 4], [26, 1, -1, 5], [27, 1, -1, 5], [28, 0, 4, 5],
    [29, 0, 4, 5], [30, 1, -1, 5], [31, -1, 4, -1], [32, 0, 5, 6],
    [33, 0, 5, 6], [34, 0, 5, 6], [35, -1, 5, -1], [36, -1, 5, -1],
    [37, 1, -1, 7], [38, 1, -1, 7], [39, 0, 6, 7], [40, 0, 6, 7]
]
df = pd.DataFrame(raw_data, columns=['word_number', 'change_type', 'doc1_paragraph_number', 'doc2_paragraph_number'])

# 调用函数生成结果
output_df = merge_word_numbers(df)
print(output_df)

代码说明 #

• para_key：针对无变更的行，用doc1和doc2的段落编号组合作为键；针对有变更的行，取有效段落编号（排除-1）作为键，确保同一段落的变更行被归为一组。
• group_id：通过对比当前行与上一行的分组键和变更类型，生成连续的分组ID，实现自动分组。
• 聚合逻辑：对每个分组筛选出对应文档中存在的词汇，同时根据分组内的变更类型集合确定最终标识，混合变更场景标记为2，完全匹配示例需求。

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