如何按指定条件拆分DataFrame中Device为B的行并分配数值

输入数据 #

初始DataFrame定义如下：

import pandas as pd

data = {
    'Device': ['A', 'B', 'C'],
    'Revenue': [1000, 2000, 5000],
    'Sales': [200, 300, 400]
}

df = pd.DataFrame(data)

实现代码 #

以下是符合需求的拆分逻辑实现：

def split_dataframe(df, x):
    # 将输入的百分比转换为小数（兼容输入为字符串如"40%"或数值0.4的情况）
    if isinstance(x, str):
        x = float(x.strip('%')) / 100
    # 当X等于40%时执行拆分（用浮点精度判断避免误差）
    if abs(x - 0.4) < 1e-9:
        # 筛选非B的行，保留原数据
        non_b_rows = df[df['Device'] != 'B'].copy()
        # 获取B行数据并拆分
        b_row = df[df['Device'] == 'B'].copy()
        # 生成B1行（40%比例）
        b1_row = b_row.assign(
            Device='B1',
            Revenue=b_row['Revenue'] * 0.4,
            Sales=b_row['Sales'] * 0.4
        )
        # 生成B2行（60%比例）
        b2_row = b_row.assign(
            Device='B2',
            Revenue=b_row['Revenue'] * 0.6,
            Sales=b_row['Sales'] * 0.6
        )
        # 合并所有行并重置索引
        result_df = pd.concat([non_b_rows, b1_row, b2_row], ignore_index=True)
        return result_df
    # 当X小于40%时直接返回原DataFrame
    elif x < 0.4:
        return df.copy()
    # X大于40%时默认返回原数据（需求未提及，此处补充兜底逻辑）
    else:
        return df.copy()

# 示例调用
# 情况1：X=40%
df_split = split_dataframe(df, 0.4)
print("X=40%时的输出：")
print(df_split)

# 情况2：X=30%
df_original = split_dataframe(df, 0.3)
print("\nX=30%时的输出：")
print(df_original)

关键逻辑说明 #

• 参数兼容：支持输入百分比字符串（如"40%"）或小数（如0.4），避免类型报错
• 精度处理：用浮点差值判断替代直接相等，解决浮点数精度问题
• 拆分逻辑：单独提取Device为B的行，按比例生成新行后与其他数据合并
• 数据安全：使用copy()避免原DataFrame被意外修改

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