如何通过布尔矩阵从Pandas DataFrame中提取对应列向量

你可以直接用pandas原生的布尔索引特性实现需求，不需要手动维护列名列表或者遍历校验，实现逻辑非常简洁：

封装调用示例 #

def get_target_df(input_index):
    # 筛选指定行值为1对应的列名
    target_cols = b_matrix.columns[b_matrix.loc[input_index] == 1]
    # 返回df对应列的结果
    return df[target_cols]

调用测试 #

输入input1时直接调用：

result = get_target_df('input1')

返回的result就是你需要的df[['col2', 'col3']]结果。

优化写法 #

如果你的b_matrix里存的0/1都是整数类型，还可以直接转布尔值简化写法：

target_cols = b_matrix.columns[b_matrix.loc[input_index].astype(bool)]

这个方法完全依托pandas内置的向量化操作实现，比手动遍历的效率更高，后续列名有变动也不需要修改代码适配，可维护性更好。

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