切片 numpy 3-d 阵列

Question

I have a 3d numpy array such as:

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]]])

Further i want to slice my array at certain random positions

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

to slice the array

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

to yield

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

What would be a faster way to perform the slicing step?

Answer 1

You can use numpy arrays instead for indexing:

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

This is documented under integer array indexing .