Python实现DataFrame同日期下C类金额匹配S类并分组标记

需求说明 #

我有一个包含date、origin、amounts字段的DataFrame，需要针对同一date，找到origin为"S"的金额——该金额等于origin为"C"的1个或多个金额之和，并将匹配的金额组用新列（如A、B、C…）标记分组。

示例数据 #

import pandas as pd

df = pd.DataFrame(
    {
        "amounts": [12, 13, 1, 14, 2, 15, 25, 29, 45],
        "date": ['2022/11/11', '2022/11/11', '2022/11/11', '2022/11/12','2022/11/12', '2022/11/12',
                 '2022/11/11', '2022/11/12', '2022/11/12'],
        "origin": ['C', 'C', 'C', 'C','C', 'S', 'S', 'S']
    }
)

示例中，2022/11/11的25（S类）等于12+13（C类），这三个属于一组；2022/11/12的29（S类）等于14+15（C类），这三个属于另一组。

现有问题 #

之前尝试按date和origin分组，用下面的函数判断是否存在匹配的组合，但这个函数只能返回布尔值，无法给匹配的金额添加分组标记：

def check_sum_exists(nums, target_sum):
    for i in range(len(nums)):
        # 单个金额匹配
        if nums[i] == target_sum:
            return True
        # 两个金额之和匹配
        for j in range(i+1, len(nums)):
            if nums[i] + nums[j] == target_sum:
                return True
        # 三个金额之和匹配
            for k in range(i+2, len(nums)):
                if nums[i] + nums[j]+ nums[k] == target_sum:
                    return True
    return False

解决方案 #

我们需要修改逻辑，不仅要找到匹配的组合，还要给每个匹配项打上分组标签。这里用回溯法来寻找C类金额的子集，匹配S类的目标金额，然后逐组标记：

import pandas as pd

def find_subsets(nums, target):
    """找到nums中能加和等于target的所有不重复子集（优先用数量少的组合）"""
    subsets = []
    nums.sort(reverse=True)  # 从大到小排序，优先匹配大金额，减少组合数

    def backtrack(start, path, current_sum):
        if current_sum == target:
            subsets.append(path.copy())
            return
        if current_sum > target:
            return
        for i in range(start, len(nums)):
            # 跳过重复值（如果有重复金额）
            if i > start and nums[i] == nums[i-1]:
                continue
            path.append(nums[i])
            backtrack(i+1, path, current_sum + nums[i])
            path.pop()

    backtrack(0, [], 0)
    # 去重（因为可能有不同顺序的同一组合），并按长度排序，优先用短组合
    unique_subsets = []
    seen = set()
    for s in subsets:
        tuple_s = tuple(sorted(s))
        if tuple_s not in seen:
            seen.add(tuple_s)
            unique_subsets.append(s)
    unique_subsets.sort(key=lambda x: len(x))
    return unique_subsets

def process_date_group(group):
    """处理单个日期分组，添加分组标记"""
    # 分离C和S类数据
    c_data = group[group['origin'] == 'C'].copy()
    s_data = group[group['origin'] == 'S'].copy()
    # 初始化分组列
    group['group'] = None
    used_c_amounts = set()  # 记录已被使用的C类金额，避免重复匹配
    group_label = ord('A')  # 分组标签从A开始

    for _, s_row in s_data.iterrows():
        target = s_row['amounts']
        # 过滤未被使用的C类金额
        available_c = c_data[~c_data['amounts'].isin(used_c_amounts)]['amounts'].tolist()
        subsets = find_subsets(available_c, target)
        if subsets:
            # 取第一个匹配的子集（优先短组合）
            matched_subset = subsets[0]
            # 标记S行的分组
            group.loc[s_row.name, 'group'] = chr(group_label)
            # 标记C行的分组
            for amount in matched_subset:
                # 找到对应的C行（注意如果有重复金额，取第一个未被标记的）
                c_indices = c_data[(c_data['amounts'] == amount) & (c_data['group'].isna())].index
                if not c_indices.empty:
                    group.loc[c_indices[0], 'group'] = chr(group_label)
                    used_c_amounts.add(amount)
            group_label += 1
    return group

# 按日期分组处理
result_df = df.groupby('date', group_keys=False).apply(process_date_group)
print(result_df)

输出结果 #

运行后会得到包含group列的DataFrame，匹配的金额会被标记为同一字母：

amounts        date origin group
0       12  2022/11/11      C     A
1       13  2022/11/11      C     A
2        1  2022/11/11      C  None
3       14  2022/11/12      C     A
4        2  2022/11/12      C  None
5       15  2022/11/12      C     A
6       25  2022/11/11      S     A
7       29  2022/11/12      S     A
8       45  2022/11/12      S  None

关键说明 #

1. find_subsets函数用回溯法寻找所有能加和等于目标的C类金额子集，并且优先返回短组合，避免一个C类金额被重复使用
2. process_date_group函数针对每个日期分组，逐个处理S类金额，标记匹配的C类金额，并更新分组标签
3. 如果有多个S类金额，会依次用A、B、C…作为分组标签，未匹配的金额group列为None

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