Python移除含至少一个0的numpy数组列的实现方法

给定原始numpy数组：

import numpy as np

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

需求：找出所有包含至少一个0的列，移除这些列并统计移除的列数，最终得到全1的3列数组。

直接用numpy的向量化操作就能高效完成，代码如下：

# 生成保留列的掩码：仅保留所有元素都不为0的列
keep_cols = (arr != 0).all(axis=0)
# 筛选得到目标数组
filtered_arr = arr[:, keep_cols]
# 计算移除的列数
removed_count = arr.shape[1] - filtered_arr.shape[1]

验证结果：

print(filtered_arr)
# 输出：
# array([[1., 1., 1.],
#        [1., 1., 1.],
#        [1., 1., 1.],
#        [1., 1., 1.]])

print(removed_count)
# 输出：2

关键逻辑说明 #

• (arr != 0).all(axis=0)：先对数组每个元素做非0判断，再沿列方向取逻辑与，得到每列是否全为非0的布尔数组，这是numpy向量化操作的典型用法，避免循环，效率极高。
• arr[:, keep_cols]：利用numpy的布尔索引直接筛选列，语法简洁且符合Pythonic风格。
• 移除列数通过原始列数与保留列数的差值计算，直观易懂。

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