pySpark 实现按目标求和条件筛选各ID下的优先匹配记录

核心逻辑匹配需求 #

• 按id分组后严格保留组内从上到下的原始行顺序
• 逐行遍历组内记录，优先锁定包含最小编号行的有效组合：
  1. 若单行val已超过组内目标sum，因所有val为正，该行不可能参与任何满足条件的组合，直接跳过
  2. 若加入当前行后累计和未超过目标值，则纳入保留集合，累计和同步更新
  3. 一旦累计和刚好等于目标sum，立即停止当前组的遍历，后续记录全部过滤
  4. 若累计和超过目标值，跳过当前行，继续检查下一行

完整可运行代码 #

import pandas as pd

# 构造示例DataFrame
df = pd.DataFrame({
    "id": [1,1,1,2,2,2],
    "val": [1,2,1,3,1,1],
    "sum_needs_to_be": [3,3,3,2,2,2]
})

def process_single_group(group_df):
    # 同一id下的目标sum固定，直接取首行值即可
    target_sum = group_df["sum_needs_to_be"].iloc[0]
    current_total = 0
    keep_row_index = []
    # 按组内原始顺序逐行遍历
    for idx, row in group_df.iterrows():
        # 累计和已达目标，直接终止遍历
        if current_total == target_sum:
            break
        row_val = row["val"]
        # 单行val超过目标、加上当前val后累计和超目标，均跳过该行
        if row_val > target_sum or current_total + row_val > target_sum:
            continue
        # 满足纳入条件，加入保留列表
        current_total += row_val
        keep_row_index.append(idx)
    # 返回过滤后的组内数据
    return group_df.loc[keep_row_index]

# 分组应用处理逻辑
result_df = df.groupby("id", group_keys=False).apply(process_single_group)

运行结果验证 #

执行上述代码后，result_df的输出完全匹配题目要求：

id  val  sum_needs_to_be
0   1    1                3
1   1    2                3
4   2    1                2
5   2    1                2

额外说明 #

• 实现逻辑严格符合“优先选择包含最小索引记录的组合”的规则，只要是包含最小编号行的有效组合，会被优先选中，不会被后续其他组合替代
• 适配场景中val为正数值的前提，性能上可直接用于大规模DataFrame的逐组遍历处理

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