如何在PySpark DataFrame中找出最接近目标值的求和元素并标记

要实现按ID分组，找出每组中能相加最接近closest_val的Weight_val子集并标记，可通过暴力枚举子集的方式实现（适合每组数据量较小的场景），具体步骤如下：

1. 导入依赖库 #

import pandas as pd
import itertools

2. 构造示例数据（或读取你的数据集） #

data = {
    'ID': [134211, 134211, 134211, 134211, 434217, 434217, 434217, 535551, 535551, 535551],
    'Weight_val': [3.9, 21, 1.5, 3.3, 15, 40, 2.1, 3.1, 7.3, 12],
    'closest_val': [5,5,5,5,19,19,19,18,18,18]
}
df = pd.DataFrame(data)

3. 定义分组处理函数 #

这个函数会针对每个ID组，枚举所有非空子集，计算子集和与目标值的差值，找到最接近目标的子集，最后为每个元素标记YES/NO：

def mark_closest_subset(group):
    weights = group['Weight_val'].values
    target = group['closest_val'].iloc[0]
    best_diff = float('inf')
    best_subset_indices = []
    
    # 枚举所有可能的非空子集
    for subset_size in range(1, len(weights) + 1):
        for indices in itertools.combinations(range(len(weights)), subset_size):
            subset_sum = sum(weights[i] for i in indices)
            current_diff = abs(subset_sum - target)
            
            # 更新最优子集：优先选差值更小的，差值相同时选和更大的
            if current_diff < best_diff:
                best_diff = current_diff
                best_subset_indices = indices
            elif current_diff == best_diff:
                current_sum = sum(weights[i] for i in indices)
                best_sum = sum(weights[i] for i in best_subset_indices)
                if current_sum > best_sum:
                    best_subset_indices = indices
    
    # 生成标记列
    group['mark'] = ['YES' if idx in best_subset_indices else 'NO' for idx in range(len(group))]
    return group

4. 分组应用函数并输出结果 #

result_df = df.groupby('ID', group_keys=False).apply(mark_closest_subset)
print(result_df)

执行后得到的结果与你期望的一致：

ID  Weight_val  closest_val mark
0  134211         3.9            5   NO
1  134211        21.0            5   NO
2  134211         1.5            5  YES
3  134211         3.3            5  YES
4  434217        15.0           19  YES
5  434217        40.0           19   NO
6  434217         2.1           19  YES
7  535551         3.1           18   NO
8  535551         7.3           18  YES
9  535551        12.0           18  YES

注意事项 #

• 暴力枚举法仅适合每组数据量较小的场景（比如每组元素数≤10），如果每组数据量很大，建议改用动态规划的近似子集和算法。
• 代码中针对差值相同的情况，选择了和更大的子集，确保结果更贴合目标值；若你有其他优先级需求，可修改对应逻辑。

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