如何更高效地对三维NumPy数组按指定位置执行切片操作？

先看我们的三维数组和目标位置列表：

import numpy as np

data = np.array([[[1, 2, 3, -999],
                  [5, 6, 7, 8],
                  [9, 10, 11, 12]],
             
                 [[10, 20, 30, -999],
                  [50, 60, 70, 80],
                  [90, 100, 110, 120]],
             
                 [[100, 200, 300, -999],
                  [500, 600, 700, 800],
                  [900, 1000, 1100, 1200]]])

pos = [[0, 1], [0, 2], [1, 0]]

原来使用列表推导式的切片方式：

slices = [data[:, p[0], p[1]] for p in pos]

可以得到结果：[[2, 20, 200], [3, 30, 300], [5, 50, 500]]

更高效的方法 #

直接利用NumPy的高级索引，通过向量化操作替代Python循环，有两种简洁实现方式：

方法一：将位置列表转为NumPy数组后索引

pos_arr = np.array(pos)
slices = data[:, pos_arr[:, 0], pos_arr[:, 1]].T

方法二：解包位置列表（代码更简洁）

slices = data[:, *zip(*pos)].T

两种方法都能得到目标结果：

array([[  2,  20, 200],
       [  3,  30, 300],
       [  5,  50, 500]])

效率优势说明 #

列表推导式本质是Python层面的逐个循环，当pos长度较大或data数组规模庞大时，循环会带来明显性能损耗。而NumPy的高级索引是底层C语言实现的向量化运算，能充分利用NumPy的优化机制，运算速度远快于Python循环，在大规模数据处理场景下优势尤为显著。

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