如何简洁筛选DataFrame中指定列均满足对应阈值条件的行子集？

给定如下DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame(data=[[0, 1, 2], 
                        [2, 3, 4], 
                        [4, 5, 6], 
                        [6, 7, 8]], 
                  columns= ["a", "b", "c"])

同时有布尔数组bool_arr指定需要应用筛选的列，以及阈值数组thresh_arr（仅bool_arr为True的列，需满足该行对应值低于阈值数组中的对应值）：

bool_arr = np.array([True, False, True])
thresh_arr = np.array([3, 2, 7])

需求是筛选出所有符合条件的行，期望结果为：

a  b  c
0  0  1  2
1  2  3  4

原实现代码如下，现寻求更简洁的写法：

sel_cols = df.columns[bool_arr]
sel_thresh = thresh_arr[bool_arr]
df_masked = df[sel_cols] < sel_thresh
sel_rows = df.index[df_masked.sum(axis=1) == bool_arr.sum()]
df_subset = df.loc[sel_rows]

简洁实现方式 #

利用pandas的广播特性和布尔索引，可以直接生成筛选掩码，一步得到目标子集：

mask = (df < thresh_arr).loc[:, bool_arr].all(axis=1)
df_subset = df[mask]

甚至可以压缩为一行：

df_subset = df[(df < thresh_arr).loc[:, bool_arr].all(axis=1)]

代码解释 #

1. df < thresh_arr：借助广播机制，让DataFrame的每一行与thresh_arr逐元素比较，生成布尔类型的DataFrame
2. .loc[:, bool_arr]：筛选出bool_arr标记为True的列对应的布尔结果
3. .all(axis=1)：检查每行是否所有指定列都满足“值低于阈值”的条件（整行全为True则保留）
4. 用最终的布尔掩码直接索引原DataFrame，得到符合要求的行子集

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