如何在Pandas中基于列除法生成列表类型新列？

原始数据 #

首先创建示例DataFrame：

import pandas as pd

df = pd.DataFrame([["X",62,5],["Y",16,3],["Z",27,4]],columns=["id","total","days"])

对应的数据表：

需求说明 #

需要新增一列plan，规则如下：

• 每个plan元素是长度等于days值的列表
• 列表默认填充total // days的商
• 如果total % days有余数，从列表末尾开始，给对应余数数量的元素各加1

预期输出 #

最终得到的DataFrame代码示例：

df_out = pd.DataFrame([["X",62,5,[12,12,12,13,13]],["Y",16,3,[5, 5, 6]],["Z",27,4,[6, 7, 7, 7]]],columns=["id","total","days","plan"])

对应数据表：

实现方案 #

方法一：自定义函数+apply #

通过apply遍历每一行，结合自定义函数生成目标列表：

import pandas as pd

def generate_plan(row):
    total = row['total']
    days = row['days']
    # 一次性算出商和余数
    quotient, remainder = divmod(total, days)
    # 初始化全为商的列表
    plan = [quotient] * days
    # 从末尾开始给余数个元素加1
    for i in range(-1, -remainder-1, -1):
        plan[i] += 1
    return plan

# 生成plan列
df['plan'] = df.apply(generate_plan, axis=1)

方法二：列表推导式简化版 #

用列表推导式直接生成列表，代码更紧凑：

def generate_plan(row):
    q, r = divmod(row['total'], row['days'])
    # 判断索引是否处于需要加1的位置（从倒数第r个元素开始）
    return [q + 1 if i >= (row['days'] - r) else q for i in range(row['days'])]

df['plan'] = df.apply(generate_plan, axis=1)

逻辑说明 #

1. divmod(a, b)：返回(a//b, a%b)，同时得到商和余数，比分开计算更高效
2. 初始列表：用[quotient] * days快速创建指定长度的列表
3. 余数处理：要么通过负索引循环修改末尾元素，要么通过索引判断直接生成带加1的列表，两种方式都能满足需求

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