Pandas中如何用向量化方法将符合条件的值替换为列名数值？

需求说明 #

我需要编写一个函数，输入是带**多级列索引（MultiIndex）**的DataFrame和布尔型掩码（比如df < 2、df >=4这类布尔DataFrame）。列索引格式为("some string", 0.4)，第二个元素是数值类型。

输出要求：

• 不符合掩码条件的值保持不变
• 符合条件的值替换为对应列名的数值部分
• 输入中的NaN值保持不变（除非掩码是df.isna()这类）

现有实现 #

我已经写了一个版本（代码文件mytable.py）：

import pandas as pd
import numpy as np


data = {
    ("A", 0.2): [4.0, 1.0, np.nan],
    ("B", 0.6): [0.0, np.nan, 4.0],
    ("C", 0.7): [0.0, 5.0, 1.0],
}
df = pd.DataFrame(data)


def replaced_with_colname(table, mask):
    series1 = (table[col][mask[col]] for col in table.columns)
    series2 = (s.apply(lambda x: s.name[1]) for s in series1)
    t2 = table.copy()
    for s in series2:
        t2.update(s)
    return t2

示例运行效果 #

$ python3 -i mytable.py

>>> df
     A    B    C
   0.2  0.6  0.7
0  4.0  0.0  0.0
1  1.0  NaN  5.0
2  NaN  4.0  1.0
>>> replaced_with_colname(df, df>2)
     A    B    C
   0.2  0.6  0.7
0  0.2  0.0  0.0
1  1.0  NaN  0.7
2  NaN  0.6  1.0

问题 #

这个实现能完成功能，但写法繁琐且可能性能不足。有没有更向量化、符合Pandas惯用风格的实现方式？也就是尽量用Pandas内置方法，减少循环。

我参考过类似问题，但场景不同：

• 有的是用表内数据替换，不是列名中的数值
• 有的是替换特定值，不是所有符合掩码条件的值

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