Python pandas按条件新增SMB列并基于分位数标记规模

原始数据集 #

import pandas as pd
import numpy as np

data = {'Year':[2012, 2013, 2012, 2013, 2014, 2013],
        'Quarter':[2, 2, 2, 2, 3, 1],
        'ID':['CH7744', 'US4652', 'CA47441', 'CH1147', 'DE7487', 'US5174'],
        'MC':[3348.22, 8542.55, 11851.2, 15718.1, 29914.7, 8731.78 ],
        'PB': [2.74, 0.95, 1.57, 2.13, 0.54, 5.32]}

df = pd.DataFrame(data)

需求说明 #

• 新增SMB列，仅处理Quarter=2的行
• 按年份+季度分组，对每组的MC列计算0.5分位数
• MC值≤分位数标记为'Small'，≥分位数标记为'Big'
• 非Quarter=2的行填充np.nan

期望输出 #

data = {'Year':[2012, 2013, 2012, 2013, 2014, 2013],
        'Quarter':[2, 2, 2, 2, 3, 1],
        'ID':['CH7744', 'US4652', 'CA47441', 'CH1147', 'DE7487', 'US5174'],
        'MC':[3348.22, 8542.55, 11851.2, 15718.1, 29914.7, 8731.78 ],
        'PB': [2.74, 0.95, 1.57, 2.13, 0.54, 5.32],
        'SMB': ['Small', 'Small', 'Big', 'Big', np.NaN, np.NaN]}

df = pd.DataFrame(data)

当前尝试代码 #

# Quantile 0.5 for MC sorting (small & big)
smbQuantile = 0.5
Years = df['Year'].unique()
dataframes_list = []

# Calculate Small and Big and merge back into dataFrame
for i in Years:
    df_temp = df.loc[(df['Year'] == i) & (df['Quarter'] == 2)]  # 修正原代码笔误：df_sb改为df
    df_temp['SMB'] = ''
    # Assign factor size based on market cap
    df_temp.SMB[df_temp.MC <= df_temp.MC.quantile(smbQuantile)] = 'Small'  # 修正原代码笔误：MKT_CAP改为MC
    df_temp.SMB[df_temp.MC >= df_temp.MC.quantile(smbQuantile)] = 'Big'
    dataframes_list.append(df_temp)

df = pd.concat(dataframes_list)

简洁实现方案 #

方案一：分组+自定义函数 #

利用Pandas的groupby和apply方法，直接在原DataFrame上完成计算，无需循环拼接：

import pandas as pd
import numpy as np

# 初始化数据
data = {'Year':[2012, 2013, 2012, 2013, 2014, 2013],
        'Quarter':[2, 2, 2, 2, 3, 1],
        'ID':['CH7744', 'US4652', 'CA47441', 'CH1147', 'DE7487', 'US5174'],
        'MC':[3348.22, 8542.55, 11851.2, 15718.1, 29914.7, 8731.78 ],
        'PB': [2.74, 0.95, 1.57, 2.13, 0.54, 5.32]}

df = pd.DataFrame(data)

# 初始化SMB列为NaN
df['SMB'] = np.nan

# 定义分组处理函数
def mark_smb(group):
    median = group['MC'].quantile(0.5)
    group['SMB'] = np.where(group['MC'] <= median, 'Small', 'Big')
    return group

# 仅处理Quarter=2的行，赋值回原DataFrame
df.loc[df['Quarter'] == 2, :] = df.loc[df['Quarter'] == 2, :].groupby(['Year', 'Quarter']).apply(mark_smb)

print(df)

方案二：嵌套np.where+transform（极简写法） #

通过groupby.transform直接获取每组的分位数，搭配np.where一步完成所有逻辑：

import pandas as pd
import numpy as np

# 初始化数据
data = {'Year':[2012, 2013, 2012, 2013, 2014, 2013],
        'Quarter':[2, 2, 2, 2, 3, 1],
        'ID':['CH7744', 'US4652', 'CA47441', 'CH1147', 'DE7487', 'US5174'],
        'MC':[3348.22, 8542.55, 11851.2, 15718.1, 29914.7, 8731.78 ],
        'PB': [2.74, 0.95, 1.57, 2.13, 0.54, 5.32]}

df = pd.DataFrame(data)

# 一步计算SMB列
df['SMB'] = np.where(
    df['Quarter'] == 2,
    np.where(
        df['MC'] <= df.groupby(['Year', 'Quarter'])['MC'].transform(lambda x: x.quantile(0.5)),
        'Small',
        'Big'
    ),
    np.nan
)

print(df)

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