如何快速获取多维NumPy数组中存在非零元素的列索引？

NumPy 数组处理方案 #

对于给定的NumPy数组，使用向量化操作可高效获取目标列索引：

import numpy as np

arr = np.array([
         [1, 0, 0.5, 1, 0],
         [0, 0, 0,   0, 0],
         [1, 0, 0,   1, 0]
        ])

# 沿行方向（axis=0）判断每列是否存在元素>0，再提取索引
target_indices = np.where(np.any(arr > 0, axis=0))[0]
print(target_indices)  # 输出: [0 2 3]

原理 #

• arr > 0：生成布尔数组，标记每个元素是否大于0
• np.any(..., axis=0)：沿列维度逐列判断是否存在至少一个True，得到每列的布尔结果
• np.where()：提取布尔结果中True对应的索引位置

Pandas DataFrame 直接处理方案 #

从DataFrame出发无需转成NumPy，直接用Pandas的向量化方法替代iterrows循环，效率提升明显：

import pandas as pd

# 修正原代码中c2的引号错误
df1 = pd.DataFrame([
                   [1, 0, 0.5, 1, 0],
                   [0, 0, 0,   0, 0],
                   [1, 0, 0,   1, 0]
                  ], 
                  columns=['c1', 'c2', 'c3', 'c4', 'c5'],
                  index=['r1', 'r2', 'r3'])

# 获取满足条件的列名
target_columns = df1.columns[df1.gt(0).any(axis=0)]
print(target_columns)  # 输出: Index(['c1', 'c3', 'c4'], dtype='object')

# 如果需要列索引，执行以下代码
target_indices = np.where(df1.gt(0).any(axis=0))[0]
print(target_indices)  # 输出: [0 2 3]

优势 #

iterrows是逐行的Python级循环，数据量越大效率越低；而gt()和any()都是Pandas底层优化的向量化操作，基于C实现，处理大规模数据时速度会快几个数量级。

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