Python中基于元素列表计算DataFrame连续元素对的求和问题

现有数据 #

Pandas DataFrame #

import pandas as pd
data = {'Name1':['A', 'B', 'C', 'D','E','F','G','H'],
        'Name2':['B','C','D','E','F','G','H','D'],
        'sum':[12,10,5,5, 10,15,10,13]
       }
df = pd.DataFrame(data)

需要计算的元素列表 #

my_list=[['A','B'],['F','G'],['A','B','C'],['A','B','C','D','E'],['E','F','G']]

计算规则 #

每个子列表的求和逻辑：

• 长度为2的子列表，直接取DataFrame中对应Name1与Name2配对的sum值
• 长度大于2的子列表，依次累加连续两两配对的sum值，例如：
  • ['A','B','C'] = ['A','B']的sum + ['B','C']的sum = 12+10=22
  • ['A','B','C','D','E'] = ['A','B']+['B','C']+['C','D']+['D','E']的sum总和=12+10+5+5=32

实现代码 #

# 构建配对到sum值的映射字典，提升查找效率
pair_sum = df.set_index(['Name1', 'Name2'])['sum'].to_dict()

# 定义计算单个列表总和的函数
def calculate_total(lst):
    total = 0
    # 遍历列表中连续的两两配对
    for i in range(len(lst)-1):
        pair = (lst[i], lst[i+1])
        total += pair_sum.get(pair, 0)  # 若配对不存在默认加0，可按需调整
    return total

# 批量计算所有子列表的结果
results = {tuple(lst): calculate_total(lst) for lst in my_list}

# 输出结果
for key, value in results.items():
    print(f"{list(key)} -> {value}")

运行结果 #

['A', 'B'] -> 12
['F', 'G'] -> 15
['A', 'B', 'C'] -> 22
['A', 'B', 'C', 'D', 'E'] -> 32
['E', 'F', 'G'] -> 25

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