如何基于多列条件生成新变量？附Python代码与数据集示例

数据集 #

变量var1生成规则 #

• 当col_dep_ever=1且regime_d与regime_o相等时，var1=2；
• 当col_dep_ever=1且regime_d与regime_o不相等时，var1=1；
• 当col_dep_ever=0时，无论regime_d与regime_o是否相等，var1=0；
• 其他情况var1=None。

待修正的Python循环代码 #

for i in range(len(???)):
    if regime_d[i]== regime_o[i] and col_dep_ever==1:
        df['var1']=2
    elif regime_d[i]!=regime_o[i] and col_dep_ever==1:
        df['var1']=1
    elif regime_d[i]!=regime_o[i] and col_dep_ever==0:
        df['var1']=0
    elif regime_d[i]==regime_o[i] and col_dep_ever==0:
        df['var1']=0
    else:
        df['var1']=None

原代码的核心问题 #

1. len(???)应填df，但循环写法本身存在逻辑错误：每次赋值df['var1']=x会覆盖整列值，而非对应行；
2. col_dep_ever未取对应行的值，直接用列判断会导致逻辑混乱；
3. pandas中循环遍历行效率极低，属于非推荐写法。

推荐写法：向量化操作（高效简洁） #

利用pandas和numpy的向量化特性，无需循环即可完成：

import numpy as np
import pandas as pd

# 构造数据集（已有df可跳过此步骤）
data = [
    ["Afghanistan", 0.0, "United Kingdom", 1.0, 1948, 1.0, 1919.0],
    ["Afghanistan", 0.0, "United Kingdom", 1.0, 1949, 1.0, 1919.0],
    ["Afghanistan", 0.0, "United Kingdom", 1.0, 1950, 1.0, 1919.0],
    ["India", 0.0, "United Kingdom", 1.0, 1948, 1.0, 1920.0],
    ["India", 0.0, "United Kingdom", 1.0, 1949, 1.0, 1920.0],
    ["Afghanistan", 0.0, "United Kingdom", 1.0, 1950, 1.0, 1920.0]
]
df = pd.DataFrame(data, columns=["country_d", "regime_d", "country_o", "regime_o", "year", "col_dep_ever", "col_dep_end_year"])

# 生成var1
df['var1'] = np.where(
    df['col_dep_ever'] == 1,
    np.where(df['regime_d'] == df['regime_o'], 2, 1),
    np.where(df['col_dep_ever'] == 0, 0, None)
)

若坚持使用循环（不推荐） #

修正后的循环代码如下，确保每行赋值正确：

# 先初始化var1列
df['var1'] = None

for i in range(len(df)):
    regime_d_val = df.loc[i, 'regime_d']
    regime_o_val = df.loc[i, 'regime_o']
    col_dep_val = df.loc[i, 'col_dep_ever']
    
    if col_dep_val == 1:
        df.loc[i, 'var1'] = 2 if regime_d_val == regime_o_val else 1
    elif col_dep_val == 0:
        df.loc[i, 'var1'] = 0
    else:
        df.loc[i, 'var1'] = None

结果验证 #

根据给定数据集，所有行col_dep_ever=1且regime_d != regime_o，最终所有行的var1值均为1，符合规则要求。

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