如何更简洁地为Pandas DataFrame反对角线下方单元格赋值NaN？

问题背景 #

现有列名从0开始递增的Pandas DataFrame，需要将左下到右上的反对角线下方的单元格值设为NaN。示例DataFrame如下：

import numpy as np
import pandas as pd

a = np.random.randint(0,10,(5,5))
df = pd.DataFrame(a, index=['A','B','C','D','E'], columns=range(5))

输出：

0   1   2   3   4
A   0   8   9   4   5
B   3   4   5   0   1
C   8   5   5   4   1
D   8   5   3   8   3
E   7   2   6   7   2

期望输出 #

0   1   2   3   4
A   0   8.0 9.0 4.0 5.0
B   3   4.0 5.0 0.0 NaN
C   8   5.0 5.0 NaN NaN
D   8   5.0 NaN NaN NaN
E   7   NaN NaN NaN NaN

现有繁琐实现 #

原方案需要新增辅助列并循环赋值，步骤冗余：

check_digit_max = df.columns[-1]
df['check_digit'] = [i for i,x in enumerate(df.index)]

for col in df.iloc[:,:-1].columns:
    df[col] = [x if int(col)+y <= check_digit_max else np.nan
               for x,y in zip(df[col], df['check_digit'])]

df = df.drop('check_digit',axis=1)

简洁解决方案 #

利用numpy广播生成布尔掩码，结合Pandas的mask()方法，一行代码即可完成：

result_df = df.mask(np.arange(len(df))[:, None] + np.arange(df.shape[1]) > df.columns[-1])

代码解释 #

1. 生成行/列索引矩阵：
   • np.arange(len(df))[:, None]：将行索引（0到4）转换为形状为(5,1)的列向量
   • np.arange(df.shape[1])：列索引（0到4）为形状(5,)的行向量
2. 广播计算行号+列号：两者相加后得到5×5的矩阵，每个元素对应单元格的行号与列号之和
3. 生成掩码：将行号+列号大于最大列号（此处为4）的位置标记为True
4. 应用掩码：df.mask(mask)会将掩码中True对应的单元格值替换为NaN

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