如何计算合同扩展关联的聚合金额与关联合同ID新列？

原始数据集 #

需求说明 #

• 当"合同扩展"值为No（不区分大小写）时，该合同为独立根合同；若值为其他合同ID，则表示当前合同是对应ID的扩展合同
• 需为每条记录生成两个字段：
  • 新合同ID：当前合同所属的所有关联合同ID（包括自身和所有上下游扩展/被扩展合同）的集合，用逗号分隔
  • 新已购聚合金额：所有关联合同的"已购数量"总和

实现代码（Python + Pandas） #

import pandas as pd

# 读取原始数据（这里模拟数据，实际可从CSV/Excel文件读取）
data = pd.DataFrame([
    ["ID 1", "A", "No", "01/01/2023", "01/05/2023", 1000],
    ["ID 1", "B", "No", "01/01/2023", "01/05/2023", 1000],
    ["ID 2", "A", "ID 1", "01/03/2023", "01/05/2023", 500],
    ["ID 2", "B", "ID 1", "01/03/2023", "01/05/2023", 500],
    ["ID 3", "A", "ID 2", "01/04/2023", "01/05/2023", 100],
    ["ID 4", "A", "NO", "01/01/2023", "01/07/2023", 2000],
    ["ID 5", "A", "ID 4", "01/03/2023", "01/07/2023", 250],
    ["ID 6", "A", "No", "01/02/2023", "01/08/2023", 1200]
], columns=["合同", "KPI", "合同扩展", "开始日期", "合同结束日期", "已购数量"])

# 统一"合同扩展"的大小写，避免大小写不一致导致判断错误
data["合同扩展"] = data["合同扩展"].str.capitalize()

# 构建唯一合同的扩展关系映射（同一合同的不同KPI记录只保留一条）
contract_map = data.drop_duplicates(subset=["合同"]).set_index("合同")["合同扩展"].to_dict()

# 定义函数：收集指定合同的所有关联合同（包含上下游扩展关系）
def get_related_contracts(contract_id):
    related = []
    current = contract_id
    # 向上遍历所有父级合同（被扩展的合同）
    while current != "No":
        related.append(current)
        current = contract_map.get(current, "No")
    # 找到根合同
    root = related[-1] if related else contract_id
    # 向下遍历所有子级合同（扩展根合同的合同）
    for c in contract_map.keys():
        if c not in related:
            temp = c
            while temp != "No":
                if temp == root:
                    related.append(c)
                    break
                temp = contract_map.get(temp, "No")
    # 去重并按ID数字排序
    related = sorted(list(set(related)), key=lambda x: int(x.split()[1]))
    return related

# 预计算每个合同的聚合结果
agg_result = {}
for contract in contract_map.keys():
    related_contracts = get_related_contracts(contract)
    total_amount = data[data["合同"].isin(related_contracts)]["已购数量"].sum()
    agg_result[contract] = {
        "新合同ID": ", ".join(related_contracts),
        "新已购聚合金额": total_amount
    }

# 将聚合结果映射回原始数据集
data["新合同ID"] = data["合同"].map(lambda x: agg_result[x]["新合同ID"])
data["新已购聚合金额"] = data["合同"].map(lambda x: agg_result[x]["新已购聚合金额"])

# 输出结果
print(data)

代码说明 #

1. 数据预处理：统一"合同扩展"字段的大小写，避免因大小写差异导致的关系判断错误
2. 合同映射构建：提取唯一合同的扩展关系，减少重复计算量
3. 关联合同收集：通过递归向上查找根合同，再遍历所有合同匹配根合同的子级扩展合同，确保收集全所有关联合同
4. 聚合与映射：按合同分组计算关联合同的已购数量总和，最后将结果映射到原始数据集的每条记录中

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