如何用Python实现按百分比分组赋值的DataFrame函数

函数实现代码 #

import pandas as pd
import numpy as np

def foo(df, perc, *args):
    total_rows = len(df)
    # 计算每个标记对应的行数，按四舍五入取整
    rows_per_segment = round(total_rows * (perc / 100))
    # 初始化标记列为空值
    df['label'] = np.nan
    
    current_start = 0
    for tag in args:
        # 计算当前区间的结束索引，避免超出DataFrame范围
        current_end = current_start + rows_per_segment
        current_end = min(current_end, total_rows)
        # 为对应区间赋值标记
        df.loc[current_start:current_end-1, 'label'] = tag
        current_start = current_end
        # 若已覆盖所有行，提前终止循环
        if current_start >= total_rows:
            break
    return df

代码说明 #

• 先将输入的百分比转换为小数，结合DataFrame总行数计算每个标记需要覆盖的行数，用round()处理非整数的情况（你也可以根据需求换成math.floor()或math.ceil()）
• 初始化新增的label列为NaN（对应问题中的null）
• 遍历传入的标记值，依次为对应行区间赋值，每次更新起始索引，直到覆盖完所有标记或所有行

使用示例 #

# 创建测试用DataFrame
population_df = pd.DataFrame({'data': range(100)})
# 调用函数，按20%区间依次标记A、B、C
result = foo(population_df, 20, 'A', 'B', 'C')
# 查看标记分布
print(result['label'].value_counts())

输出结果：

A    20
B    20
C    20
Name: label, dtype: int64

剩余40行的label值为NaN

注意事项 #

• 函数默认会直接修改传入的原DataFrame，如果需要保留原数据，可以在函数开头添加df = df.copy()
• 若传入的百分比乘以总行数后为非整数，四舍五入可能导致最后一个标记的行数略有偏差，可根据业务需求调整取整方式

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