如何使用numpy根据"dataframe"中的"key"列值来映射多个列？

Question

运行这段代码：

import pandas as pd
import numpy as np

# 创建一个示例 DataFrame
data = {
    'key': ['key1', 'key2', 'key3', 'key1', 'key2'],
    'colA': ['value1A', 'value2A', 'value3A', 'value4A', 'value5A'],
    'colB': ['value1B', 'value2B', 'value3B', 'value4B', 'value5B'],
    'colC': ['value1C', 'value2C', 'value3C', 'value4C', 'value5C'],
    'colD': ['value1D', 'value2D', 'value3D', 'value4D', 'value5D']
}

df = pd.DataFrame(data)

# 创建条件
conditions = [df['key'] == 'key1',
            df['key'] == 'key2',
            df['key'] == 'key3']

# 根据条件对相应列应用选择
df['colA'] = np.select([df['key'] == 'key1'], [df['colA']], default='NA')
df['colD'] = np.select([df['key'] == 'key1'], [df['colD']], default='NA')
df['colB'] = np.select([df['key'] == 'key2'], [df['colB']], default='NA')
df['colC'] = np.select([df['key'] == 'key3'], [df['colC']], default='NA')

# 显示结果 DataFrame
print(df)

此段代码生成以下输出：

         key     colA     colB     colC     colD
0  key1  value1A       NA       NA  value1D
1  key2       NA  value2B       NA       NA
2  key3       NA       NA  value3C       NA
3  key1  value4A       NA       NA  value4D
4  key2       NA  value5B       NA       NA

问题：是否有更有效率的方法重写这段代码，以便我不必为每一列都执行numpy.select操作？在本例中，“key”列实质上决定了该行哪些列包含有效数据。如果数据无效，我想将其标记为NA；如果有效，则保留行中的原始值。在我的实际数据集中，“key”列控制了如何映射多列（例如key1控制了colA和colD的映射）。我倾向于使用numpy或其他向量化方法，因为它们通常比其他方法如map更快，但我愿意听取所有建议。

Stephen Cleary · Answer

你可以采用向量化的方式构建一个基于键/列的字典，并通过pandas.get_dummies获取虚拟变量，然后利用这个虚拟变量矩阵配合pandas.DataFrame.mask方法来屏蔽不需要的值：

d = {'key1': ['colA', 'colD'],
     'key2': ['colB'],
     'key3': ['colC'],
    }

s = pd.Series(d).explode()
mask = pd.get_dummies(s, dtype=bool).groupby(level=0).max()

# 获取除'key'之外的所有列
cols = df.columns.difference(['key'])

# 使用mask对数据框中相应的列应用条件
df[cols] = df[cols].where(mask.reindex(df['key']).to_numpy(), 'NA')

输出结果：

    key     colA     colB     colC     colD
0  key1  value1A       NA       NA  value1D
1  key2       NA  value2B       NA       NA
2  key3       NA       NA  value3C       NA
3  key1  value4A       NA       NA  value4D
4  key2       NA  value5B       NA       NA

中间步骤：

# 构建的mask
       colA   colB   colC   colD
key1   True  False  False   True
key2  False   True  False  False
key3  False  False   True  False

# mask.reindex(df['key']).to_numpy()得到的数组
[[ True False False  True]
 [False  True False False]
 [False False  True False]
 [ True False False  True]
 [False  True False False]]

另一种方法可能是通过对数据重塑并使用pandas.DataFrame.merge过滤：

d = {'key1': ['colA', 'colD'],
     'key2': ['colB'],
     'key3': ['colC'],
    }

(df.reset_index()
   .melt(id_vars=['index', 'key'])
   .merge(pd.Series(d).explode().rename_axis('key').reset_index(name='variable'))
   .set_index(['index', 'key', 'variable'])['value']
   .unstack('variable', fill_value='NA')
   .reset_index('key').rename_axis(index=None, columns=None)
)

haccks · Answer

理解您的需求后，您可以使用pandas.DataFrame.where来替换那些不符合正确键/列组合的位置。我将这些组合存储在一个字典中，并遍历它们。尽管循环不是非常高效，但如果您的列和键数量不多，这种方法应该足够快并且更重要的是易于理解。

mappings = {"key1": ["colA", "colD"], "key2": ["colB"], "key3": ["colC"]}
for key, cols in mappings.items():
    for col in cols:
        df[col] = df.where(df["key"] == key)[col]

执行上述代码后，您会得到如下结果：

    key     colA     colB     colC     colD
0  key1  value1A      NaN      NaN  value1D
1  key2      NaN  value2B      NaN      NaN
2  key3      NaN      NaN  value3C      NaN
3  key1  value4A      NaN      NaN  value4D
4  key2      NaN  value5B      NaN      NaN

若需要将NaN替换成"NA"，则可以在df.where调用中添加这个参数：

mappings = {"key1": ["colA", "colD"], "key2": ["colB"], "key3": ["colC"]}
for key, cols in mappings.items():
    for col in cols:
        df[col] = df.where(df["key"] == key, "NA")[col]

执行以上修改后的代码，将得到如下结果：

    key     colA     colB     colC     colD
0  key1  value1A       NA       NA  value1D
1  key2       NA  value2B       NA       NA
2  key3       NA       NA  value3C       NA
3  key1  value4A       NA       NA  value4D
4  key2       NA  value5B       NA       NA

Greg · Answer

迭代处理列其实可以较为高效，因为在整个列上操作，而非逐行。最终，您需要根据您的具体数据进行测试，特别是随着数据规模的增长：

mapping = {'key1': ['colA', 'colD'], 'key2': 'colB', 'key3': 'colC'}
new_df = df.copy()
for key, value in mapping.items():
    selected_columns = df.loc[:, value]
    masked_values = selected_columns.where(df['key'] == key, 'NA')
    new_df.loc[:, value] = masked_values

print(new_df)

执行这段代码后，将会得到以下结果：

    key     colA     colB     colC     colD
0  key1  value1A       NA       NA  value1D
1  key2       NA  value2B       NA       NA
2  key3       NA       NA  value3C       NA
3  key1  value4A       NA       NA  value4D
4  key2       NA  value5B       NA       NA