求助：如何对NumPy数组实现‘取二跳二’的通用处理？

问题示例 #

输入数组：

import numpy as np

input_arr = np.array([[[ 1.,  0.,  2.,  0.,  3.,  0.,  4.,  0.,  5.],
        [ 6.,  0.,  7.,  0.,  8.,  0.,  9.,  0., 10.],
        [11.,  0., 12.,  0., 13.,  0., 14.,  0., 15.]],

       [[16.,  0., 17.,  0., 18.,  0., 19.,  0., 20.],
        [21.,  0., 22.,  0., 23.,  0., 24.,  0., 25.],
        [26.,  0., 27.,  0., 28.,  0., 29.,  0., 30.]]])

期望输出：

array([[[ 1.,  0.,  3.,  0.,  5.],
        [ 6.,  0.,  8.,  0., 10.],
        [11.,  0., 13.,  0., 15.]],

       [[16.,  0., 18.,  0., 20.],
        [21.,  0., 23.,  0., 25.],
        [26.,  0., 28.,  0., 30.]]])

通用解决方案 #

核心思路是按每2列为一组，保留偶数序号的组（第0、2、4...组），该方法不受数组列数、维度层级限制（只要最后一维为列维度），可灵活调整适配类似分组筛选需求。

实现步骤 #

• 生成列索引掩码：对每个列索引i，判断其所属的双列组序号i//2是否为偶数
• 通过掩码对数组最后一维（列）进行筛选

代码实现 #

# 生成列筛选掩码
col_mask = (np.arange(input_arr.shape[-1]) // 2) % 2 == 0
# 应用掩码得到目标数组
output_arr = input_arr[..., col_mask]

print(output_arr)

逻辑解释 #

1. np.arange(input_arr.shape[-1])：生成对应列数的索引数组（示例中为[0,1,2,3,4,5,6,7,8]）
2. i // 2：将索引按每2个一组划分，得到组号数组（示例中为[0,0,1,1,2,2,3,3,4]）
3. (组号) % 2 == 0：筛选出属于偶数组的索引，得到掩码[True,True,False,False,True,True,False,False,True]
4. input_arr[..., col_mask]：通过掩码对数组最后一维进行索引，保留目标列

若需调整分组大小（比如每3列为一组），只需将// 2替换为// N（N为每组列数），同时修改判断条件即可，通用性极强。

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