如何在DataFrame中按相同ID匹配日期并生成对应模块编号列？

我有一个DataFrame，想要判断每行的date是否与同ID行的module1_date至module5_date中的某一值匹配：匹配则返回对应模块编号，不匹配则返回0，生成类似示例中高亮的new列。但我当前尝试的函数会跨ID进行匹配，不符合需求，请问该如何实现？

尝试的代码 #

def function(x):
    if x in df.module1_date.values:
        return 1
    if x in df.module2_date.values:
        return 2
    if x in df.module3_date.values:
        return 3
    if x in df.module4_date.values:
        return 4
    if x in df.module5_date.values:
        return 5
    else:
        return 0

解决方案 #

你的代码问题在于每次判断时都调用了整个DataFrame的moduleN_date列所有值，自然会跨ID匹配。要实现同ID内的匹配，需要针对每行，只检查当前行对应的module1_date到module5_date的值。

方法1：行级判断（直观易读） #

直接在每行内对比date与当前行的模块日期列，返回对应编号：

def check_match(row):
    target_date = row['date']
    # 遍历module1到module5列
    for module_num in range(1, 6):
        col_name = f'module{module_num}_date'
        if row[col_name] == target_date:
            return module_num
    return 0

# 生成new列
df['new'] = df.apply(check_match, axis=1)

方法2：向量化操作（大数据量更高效） #

如果你的DataFrame数据量较大，向量化操作比apply的行级循环效率更高：

# 建立模块列名与编号的映射
module_mapping = {f'module{i}_date': i for i in range(1, 6)}

# 生成每个模块列与date列的布尔匹配结果
match_results = df.assign(**{col: df[col] == df['date'] for col in module_mapping})

# 提取每行第一个匹配的模块编号，无匹配则返回0
df['new'] = match_results[module_mapping.keys()].idxmax(axis=1).map(module_mapping).fillna(0).astype(int)

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