来自numpy ndarray的有效样本向量

Question

I have a multidimensional numpy array X of shape : (B, dim, H, W) I would like to randomly sample k dim -dimensional vectors out of X .
I can get the sample indices from a msk of shape (B, 1, H, W) :

sIdx = random.sample((msk.flat>=0).nonzero()[0], k) 

Equivalent sampling code using numpy is:

sIdx = np.random.choice((msk.flat>=0).nonzero()[0], replace=False, size=(k,))

But how can I efficiently slice X according to the "flat" sampled indices sIdx ?
That is, is there an efficient way to combine the random sampling of msk with the slicing of X ?

Answer 1

Get the respective indices for the rest of those three axes with np.unravel_index from the flattened indices and simply index into X along those axes for the final output, like so -

I,J,K = np.unravel_index(sIdx, (B, H, W))
out = X[I,:,J,K]