如何在Python中基于列条件生成BrandNewCateg_A列

现有一份销售数据集，记录了客户在A、B两类商品间的金额变动。已通过以下Python代码生成了TotalPerCxCateg（客户月度分类总金额）和TotalPerCx（客户月度总金额）两列：

df['TotalPerCxCateg'] = df.groupby(['Date','Client', 'Categ'])['Amount'].transform('sum')
df['TotalPerCx'] = df.groupby(['Date','Client'])['Amount'].transform('sum')

原始数据集 #

Client   Date     ProdID    Amount  Categ     
ABC Inc  2022-06      4      50      A           
ABC Inc  2022-06      5      50      B           
ABC Inc  2022-07      4     120      A           
ABC Inc  2022-08      4      40      A           
ABC Inc  2022-08      5      40      B           
ABC Inc  2022-08      6      30      B           

需求 #

需要新增一列BrandNewCateg_A，满足以下规则：

• 仅当客户当月总金额（TotalPerCx）较上月增长，且A类商品当月总金额（TotalPerCxCateg）较上月增长时，填入当月总金额的新增值（例如2022年7月该列值为20，因TotalPerCx较上月增长20，A类TotalPerCxCateg增长70）
• 不满足条件时填0

期望输出 #

Client   Date   ProdID  Amount  Categ TotalPerCxCateg TotalPerCx BrandNewCateg_A
ABC Inc  Jun-22   4      50      A           50             100       0 
ABC Inc  Jun-22   5      50      B           50             100       0
ABC Inc  Jul-22   4     120      A           120            120       20
ABC Inc  Aug-22   4      40      A           40             110       0
ABC Inc  Aug-22   5      40      B           70             110       0
ABC Inc  Aug-22   6      30      B           70             110       0

按以下步骤实现需求：

1. 转换日期格式为datetime类型，方便按月计算环比
2. 提取每个客户每月的总金额和A类分类金额
3. 用shift()函数获取上月数据，计算两类金额的增量
4. 根据条件判断赋值，最后将结果合并回原数据集

具体代码如下：

import pandas as pd

# 构造原始数据集
data = {
    'Client': ['ABC Inc']*6,
    'Date': ['2022-06', '2022-06', '2022-07', '2022-08', '2022-08', '2022-08'],
    'ProdID': [4,5,4,4,5,6],
    'Amount': [50,50,120,40,40,30],
    'Categ': ['A','B','A','A','B','B']
}
df = pd.DataFrame(data)

# 生成需求中已有的聚合列
df['TotalPerCxCateg'] = df.groupby(['Date','Client', 'Categ'])['Amount'].transform('sum')
df['TotalPerCx'] = df.groupby(['Date','Client'])['Amount'].transform('sum')

# 转换日期格式并分组提取月度核心数据
df['Date'] = pd.to_datetime(df['Date'])
monthly_agg = df.groupby(['Client', 'Date']).agg(
    TotalPerCx=('TotalPerCx', 'first'),
    A_Total=('TotalPerCxCateg', lambda x: x[df.loc[x.index, 'Categ'] == 'A'].iloc[0] if any(df.loc[x.index, 'Categ'] == 'A') else 0)
).reset_index()

# 计算上月的对应数据
monthly_agg['Prev_TotalPerCx'] = monthly_agg.groupby('Client')['TotalPerCx'].shift(1)
monthly_agg['Prev_A_Total'] = monthly_agg.groupby('Client')['A_Total'].shift(1)

# 计算增量并生成目标列
monthly_agg['Total_Growth'] = monthly_agg['TotalPerCx'] - monthly_agg['Prev_TotalPerCx']
monthly_agg['A_Growth'] = monthly_agg['A_Total'] - monthly_agg['Prev_A_Total']
monthly_agg['BrandNewCateg_A'] = monthly_agg.apply(
    lambda row: row['Total_Growth'] if (row['Total_Growth'] > 0 and row['A_Growth'] > 0) else 0,
    axis=1
).fillna(0)

# 合并回原数据集并格式化日期
df = df.merge(monthly_agg[['Client', 'Date', 'BrandNewCateg_A']], on=['Client', 'Date'], how='left')
df['Date'] = df['Date'].dt.strftime('%b-%y')

# 输出结果
print(df.to_string(index=False))

运行代码后即可得到符合要求的数据集。

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