多条件拆分DataFrame为无属性重叠数据集时行为异常排查

核心需求 #

将DataFrame拆分为setA和setB，满足两个条件：

1. 对于每个属性列（如prop1/prop2/prop3），setA中该列的所有值与setB中该列的所有值完全无重叠；
2. setA的规模接近指定阈值（如原数据集的70%）。

原代码错误定位 #

你的代码逻辑存在根本性缺陷：每次仅移除与当前处理行共享属性的行，但未跟踪setA中已积累的所有属性值，导致后续留在setB的行可能与setA中更早加入的行存在属性值重叠。

举个饮食数据的例子：

• 第一次处理person1（veggies=paprika, fruits=apple, meats=chicken），会移除所有veggies=paprika、fruits=apple或meats=chicken的行（包括person1-4、person10）；
• 但person9的veggies=potato，fruits=cococonut，meats=cow，与person1无任何属性重叠，因此留在原df中；
• 而setA中已有person2-3的veggies=potato，最终setB的person9的veggies值与setA重叠，违反需求。

修正方案 #

正确的逻辑需要跟踪setA已使用的所有属性值，确保后续加入setA的行或留在setB的行，其属性值均不与setA已有的属性值重叠。

修正后的代码 #

import pandas as pd
import string
import random

# 生成测试数据
names = [''.join(random.choices(string.ascii_uppercase + string.digits, k=5)) for i in range(100)]
prop1 = [random.randint(0, 100) for i in range(100)]
prop2 = [random.randint(0, 100) for i in range(100)]
prop3 = [random.randint(0, 100) for i in range(100)]

df = pd.DataFrame({'name': names, 'prop1': prop1, 'prop2': prop2, 'prop3': prop3})

# 初始化集合与已使用属性值
setA = pd.DataFrame(columns=df.columns)
setB = pd.DataFrame(columns=df.columns)
target_size = len(df) * 0.7

# 跟踪每个属性列已使用的值
used = {
    'prop1': set(),
    'prop2': set(),
    'prop3': set()
}

# 先打乱数据集，确保随机选取
shuffled_df = df.sample(frac=1, random_state=42).reset_index(drop=True)

while len(setA) < target_size and len(shuffled_df) > 0:
    # 找到第一行所有属性值都未被使用的行
    valid_rows = shuffled_df[
        (~shuffled_df['prop1'].isin(used['prop1'])) &
        (~shuffled_df['prop2'].isin(used['prop2'])) &
        (~shuffled_df['prop3'].isin(used['prop3']))
    ]
    
    if valid_rows.empty:
        break  # 没有符合条件的行，提前终止
    
    # 取第一行有效行
    selected_row = valid_rows.iloc[0]
    p1, p2, p3 = selected_row['prop1'], selected_row['prop2'], selected_row['prop3']
    
    # 筛选所有包含p1/p2/p3的行（这些行将加入setA）
    to_add = shuffled_df[
        (shuffled_df['prop1'] == p1) |
        (shuffled_df['prop2'] == p2) |
        (shuffled_df['prop3'] == p3)
    ]
    
    # 更新已使用属性值
    used['prop1'].update(to_add['prop1'].unique())
    used['prop2'].update(to_add['prop2'].unique())
    used['prop3'].update(to_add['prop3'].unique())
    
    # 将行加入setA，并从原df移除
    setA = pd.concat([setA, to_add], ignore_index=True)
    shuffled_df = shuffled_df.drop(to_add.index).reset_index(drop=True)

# 剩余行归入setB
setB = shuffled_df

# 验证条件
assert setB['prop1'].isin(used['prop1']).sum() == 0, "setB存在与setA重叠的prop1值"
assert setB['prop2'].isin(used['prop2']).sum() == 0, "setB存在与setA重叠的prop2值"
assert setB['prop3'].isin(used['prop3']).sum() == 0, "setB存在与setA重叠的prop3值"
print(f"setA规模: {len(setA)}, 占比: {len(setA)/len(df):.2f}")
print(f"setB规模: {len(setB)}, 占比: {len(setB)/len(df):.2f}")

代码说明 #

1. 打乱数据集：确保每次选取的行是随机的，避免固定顺序影响结果；
2. 跟踪已使用属性值：用used字典记录setA中每个属性列的所有值，确保后续选取的行不会与这些值重叠；
3. 筛选有效行：每次仅选取所有属性值均未被使用的行，避免引入与setA重叠的属性；
4. 批量处理关联行：将所有与选中行共享任意属性的行加入setA，并更新已使用属性集合；
5. 验证条件：最后通过断言确保setB与setA的属性值完全无重叠。

补充示例验证（饮食数据） #

将上述代码应用到饮食数据时，会自动划分出属性完全无重叠的setA和setB。例如：

• setA可能包含所有fruits=apple的行（person1-4），其veggies值为paprika、potato，meats值为chicken、cow；
• setB则包含剩余行，其fruits值为banana、cococonut，veggies值为yam、zuchinni，meats值为sheep、fish，与setA的属性值完全无重叠。

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