Pandas中基于TYPE分组计算指定列求和并新增列问题

我编写了以下Pandas代码，尝试基于EXPIRYDT和TYPE字段分组，计算指定列的求和值并新增列：

data = {'SYMBOL': ['AAAA','AAAA','AAAA','AAAA','AAAA','AAAA','AAAA'] ,
    'EXPIRYDT': ['26-Oct-23','26-Oct-23','26-Oct-23','26-Oct-23','26-Oct-23','26-Oct-23','26-Oct-23'], 
    'STRIKE': [480, 500, 525, 425, 450, 480, 500],
    'TYPE': ['CE', 'CE', 'CE', 'PE', 'PE', 'PE', 'PE'],
    'CONTRACTS': [1, 31, 1, 0, 12, 2, 6],
    'OPENINT': [4000, 25000, 1000, 1000, 64000, 2000, 5000],
    'TIMESTAMP': ['4-Sep-23','4-Sep-23','4-Sep-23','4-Sep-23','4-Sep-23','4-Sep-23','4-Sep-23']}

df=pd.DataFrame(data)
result = df.groupby(['EXPIRYDT', 'TYPE'])

df['CE_CONT'] = result['CONTRACTS'].transform('sum')
df['PE_CONT'] = result['CONTRACTS'].transform('sum')
df['CE_OI'] = result['OPENINT'].transform('sum')
df['PE_OI'] = result['OPENINT'].transform('sum')

print(df)

但未得到预期输出，我需要的预期输出如下：

SYMBOL  EXPIRYDT  STRIKE TYPE CONTRACTS  OPENINT    TIMESTAMP  CE_CONT PE_CONT CE_OI PE_OI 
AAAA    26-Oct-23   480    CE      1        4000     4-Sep-23    33     20     30000  72000
AAAA    26-Oct-23   500    CE      31       25000    4-Sep-23    33     20     30000  72000
AAAA    26-Oct-23   525    CE      1        1000     4-Sep-23    33     20     30000  72000
AAAA    26-Oct-23   425    PE      0        1000     4-Sep-23    33     20     30000  72000
AAAA    26-Oct-23   450    PE     12        64000    4-Sep-23    33     20     30000  72000
AAAA    26-Oct-23   480    PE      2        2000     4-Sep-23    33     20     30000  72000
AAAA    26-Oct-23   500    PE      6        5000     4-Sep-23    33     20     30000  72000

分组后需实现：

• 将TYPE为CE的OPENINT求和结果存入CE_OI列
• 将TYPE为PE的OPENINT求和结果存入PE_OI列
• 将TYPE为CE的CONTRACTS求和结果存入CE_CONT列
• 将TYPE为PE的CONTRACTS求和结果存入PE_CONT列

问题原因 #

你之前的代码按['EXPIRYDT', 'TYPE']分组后，transform('sum')只会返回当前分组的求和值——CE组的行只能拿到CE的和，PE组的行只能拿到PE的和，没法把两类的结果都放到每一行里。要实现需求，得换种分组方式或者处理逻辑。

解决方案1：分组求和后合并 #

先按EXPIRYDT和TYPE分组计算总和，再将TYPE展开为列名，最后合并回原表：

import pandas as pd

data = {'SYMBOL': ['AAAA','AAAA','AAAA','AAAA','AAAA','AAAA','AAAA'] ,
    'EXPIRYDT': ['26-Oct-23','26-Oct-23','26-Oct-23','26-Oct-23','26-Oct-23','26-Oct-23','26-Oct-23'], 
    'STRIKE': [480, 500, 525, 425, 450, 480, 500],
    'TYPE': ['CE', 'CE', 'CE', 'PE', 'PE', 'PE', 'PE'],
    'CONTRACTS': [1, 31, 1, 0, 12, 2, 6],
    'OPENINT': [4000, 25000, 1000, 1000, 64000, 2000, 5000],
    'TIMESTAMP': ['4-Sep-23','4-Sep-23','4-Sep-23','4-Sep-23','4-Sep-23','4-Sep-23','4-Sep-23']}

df = pd.DataFrame(data)

# 分组求和并将TYPE转为列
sum_df = df.groupby(['EXPIRYDT', 'TYPE'])[['CONTRACTS', 'OPENINT']].sum().unstack()
# 重命名列名匹配需求
sum_df.columns = [f'{col[1]}_{col[0]}' for col in sum_df.columns]
# 合并回原表
df = df.merge(sum_df, on='EXPIRYDT', how='left')
# 调整列名和顺序
df.rename(columns={
    'CE_CONTRACTS':'CE_CONT', 
    'PE_CONTRACTS':'PE_CONT', 
    'CE_OPENINT':'CE_OI', 
    'PE_OPENINT':'PE_OI'
}, inplace=True)
df = df[['SYMBOL', 'EXPIRYDT', 'STRIKE', 'TYPE', 'CONTRACTS', 'OPENINT', 'TIMESTAMP', 'CE_CONT', 'PE_CONT', 'CE_OI', 'PE_OI']]

print(df)

解决方案2：transform结合条件判断 #

按EXPIRYDT分组后，用transform配合布尔索引计算对应TYPE的总和：

import pandas as pd

data = {'SYMBOL': ['AAAA','AAAA','AAAA','AAAA','AAAA','AAAA','AAAA'] ,
    'EXPIRYDT': ['26-Oct-23','26-Oct-23','26-Oct-23','26-Oct-23','26-Oct-23','26-Oct-23','26-Oct-23'], 
    'STRIKE': [480, 500, 525, 425, 450, 480, 500],
    'TYPE': ['CE', 'CE', 'CE', 'PE', 'PE', 'PE', 'PE'],
    'CONTRACTS': [1, 31, 1, 0, 12, 2, 6],
    'OPENINT': [4000, 25000, 1000, 1000, 64000, 2000, 5000],
    'TIMESTAMP': ['4-Sep-23','4-Sep-23','4-Sep-23','4-Sep-23','4-Sep-23','4-Sep-23','4-Sep-23']}

df = pd.DataFrame(data)

# 计算CE和PE的CONTRACTS总和
df['CE_CONT'] = df.groupby('EXPIRYDT')['CONTRACTS'].transform(lambda x: x[df.loc[x.index, 'TYPE'] == 'CE'].sum())
df['PE_CONT'] = df.groupby('EXPIRYDT')['CONTRACTS'].transform(lambda x: x[df.loc[x.index, 'TYPE'] == 'PE'].sum())

# 计算CE和PE的OPENINT总和
df['CE_OI'] = df.groupby('EXPIRYDT')['OPENINT'].transform(lambda x: x[df.loc[x.index, 'TYPE'] == 'CE'].sum())
df['PE_OI'] = df.groupby('EXPIRYDT')['OPENINT'].transform(lambda x: x[df.loc[x.index, 'TYPE'] == 'PE'].sum())

print(df)

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