如何用for循环实现销售数据分层购买列计算，全列为空时终止

首先给出原始的月度销售DataFrame：

import pandas as pd
import numpy as np

data = [['2021-09-01', 100], ['2021-10-01', 200], ['2021-11-01', 300],
           ['2021-12-01', 400], ['2022-01-01', 500], ['2022-02-01', 600],
           ['2022-03-01', 700], ['2022-04-01', 800], ['2022-05-01', 900],
           ['2022-06-01', 1000], ['2022-07-01', 1100], ['2022-08-01', 1200],
           ['2022-09-01', 1300], ['2022-10-01', 1400], ['2022-11-01', 1500],
           ['2022-12-01', 1600], ['2023-01-01', 1700], ['2023-02-01', 1800],
           ['2023-03-01', 1900], ['2023-04-01', 2000], ['2023-05-01', 2100],
           ['2023-06-01', 2200], ['2023-07-01', 2300], ['2023-08-01', 2400]] 

df = pd.DataFrame(data, columns = ['Month Year', 'Sales'])

需求是按以下规则生成各次购买列：

• 初始购买列：当df['Sales'] - df['Sales'].shift(3)*factor为NaN时，取值为Sales，否则取(Sales - Sales.shift(3))*factor
• 2nd Purchase：初始购买列移位3个月后乘以因子0.9
• 3rd及之后的购买列：前一次购买列移位6个月后乘以因子0.9
• 要求用循环实现，当新增的列全为NaN时自动停止循环，替代手动逐个添加列的繁琐操作

步骤说明： #

1. 先手动计算初始购买列，因为它的逻辑和后续复购列不同
2. 初始化循环的关键参数：当前列名、移位周期（第一次复购用3，之后固定用6）、购买次数计数
3. 进入循环，每次基于前一列生成新的购买列，然后检查新列是否全为NaN，若是则终止循环

完整代码： #

factor = 0.9

# 1. 先计算初始购买列
df['Initial Purchase'] = np.where(
    pd.isnull(df['Sales'] - df['Sales'].shift(3) * factor),
    df['Sales'],
    (df['Sales'] - df['Sales'].shift(3)) * factor
)

# 2. 初始化循环参数
current_col = 'Initial Purchase'
# 第一次复购移位3个月，之后都是6个月
shift_periods = [3, 6]
purchase_num = 2  # 下一个要生成的是2nd Purchase

while True:
    # 确定当前要用的移位周期：第一次用3，之后固定用6
    shift = shift_periods[0] if purchase_num == 2 else shift_periods[1]
    # 生成新列名
    new_col = f'{purchase_num}th Purchase'
    # 计算新列值
    df[new_col] = df[current_col].shift(shift) * factor
    
    # 检查新列是否全为NaN，是的话就删除该列并终止循环
    if df[new_col].isna().all():
        df.drop(columns=[new_col], inplace=True)
        break
    
    # 更新循环变量，准备下一次迭代
    current_col = new_col
    purchase_num += 1

代码解释： #

• 初始购买列的计算逻辑和手动方式完全一致，确保结果对齐
• 循环里通过purchase_num判断当前是第一次复购还是后续复购，从而选择对应的移位周期（3或6）
• 每次生成新列后，用isna().all()检查是否全为NaN，如果是就删掉该列并停止循环，避免生成无效列
• 循环自动更新当前列名和购买次数，直到无法生成有效列为止

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