python - Pandas Dataframe 列平均值的 bool 运算 - 这必须很简单

Question

我有 pandas 数据框:

import pandas as pd
import numpy as np

df = pd.DataFrame({'A':[1,5,3],
                  'B': [4,2,6]})

df['avg'] = df.mean(axis=1)
df[df<df['avg']]

我想保留数据框中低于 df['avg'] 列平均值的所有值。当我执行以下操作时，我将返回所有 NAN

df[df<df['avg']]

如果我设置一个 for 循环，我可以获得我想要的 bool 值。

col_names = ['A', 'B']
for colname in col_names:
    df[colname] = df[colname]<df['avg']

我正在搜索的内容如下所示:

df_desired = pd.DataFrame({
    'A':[1,np.nan,3],
    'B':[np.nan,2,np.nan],
    'avg' :[2.5, 3.5, 4.5]
})

我该怎么做？必须有一种Python式的方法来做到这一点。

Answer 1

您可以使用.mask(..) [pandas-doc]这里。我们可以使用 numpy 的广播来生成高于给定平均值的 bool 值数组:

>>> df.mask(df.values > df['avg'].values[:,None])
     A    B  avg
0  1.0  NaN  2.5
1  NaN  2.0  3.5
2  3.0  NaN  4.5