如何用Janitor或其他方法实现同行列条件判断生成新列？

我需要实现以下逻辑：对DataFrame的每一行，若colX的内容包含在colZ的内容中，新增的result列取值为colZ；否则取值为colA。我试过用janitor的case_when方法，分别用str.contains和isin做判断，但没得到预期结果。实际数据更复杂，但下面的示例足以说明问题，非janitor的解决方案也可以接受。

示例数据 #

import pandas as pd

df = pd.DataFrame(
    {
        'colZ': ["zang", "zang", "zang", "z", "zang"],
        'colX': ["A", "B", "B", "A", "Z"],
        'colA': ["1", "1", "1", "1", "1"],
    }
)

期望输出 #

output_df = pd.DataFrame(
    {
        'colZ': ["zang", "zang", "zang", "z", "zang"],
        'colX': ["A", "B", "B", "A", "Z"],
        'colA': ["1", "1", "1", "1", "1"],
        'result': ["zang", "1", "1", "1", "zang"]
    }
)

我尝试过的代码 #

import janitor as jn

# 尝试用str.contains
output_df = jn.case_when(df,
                  df['colZ'].str.contains(df['colX']),  df['colZ'],
                  df['colA'],
                  column_name='result')

# 尝试用isin
output_df = jn.case_when(df,
                  df['colZ'].isin(df['colX']),  df['colZ'],
                  df['colA'],
                  column_name='result')

问题原因 #

之前的代码逻辑错误：

• df['colZ'].str.contains(df['colX'])：该写法是将colX的所有值作为匹配模式，批量匹配colZ的每个元素，并非逐行判断colX[i]是否包含在colZ[i]中
• df['colZ'].isin(df['colX'])：该写法是判断colZ的元素是否存在于colX的整个数组中，和逐行的字符串包含关系无关

解决方案 #

方法1：numpy.where + 逐行判断 #

用apply逐行检查包含关系，配合numpy.where生成结果列：

import numpy as np

df['result'] = np.where(
    df.apply(lambda row: row['colX'] in row['colZ'], axis=1),
    df['colZ'],
    df['colA']
)

方法2：janitor的case_when（修正条件） #

将逐行判断的逻辑作为case_when的条件：

output_df = jn.case_when(
    df,
    df.apply(lambda row: row['colX'] in row['colZ'], axis=1), df['colZ'],
    df['colA'],
    column_name='result'
)

方法3：numpy.select（多条件场景更灵活） #

如果后续有更多判断条件，np.select扩展性更强：

conditions = [df.apply(lambda row: row['colX'] in row['colZ'], axis=1)]
choices = [df['colZ']]
df['result'] = np.select(conditions, choices, default=df['colA'])

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