使用另一3D指數陣列對4D陣列進行索引

Question

A具有4D陣列M（a x b x c x d）和指數（3 x f）的陣列I，例如，

I = np.array([1,2,3, ...], [2,1,3, ...], [4,1,6, ...]) 

我想用I以矩陣X有f行和d列，其中到達：

X[0,:] = M[1,2,4,:] 
X[1,:] = M[2,1,1,:] 
X[2,:] = M[3,3,6,:] 
... 

我知道我可以使用M[I[0], I[1], I[2]]，但是，我想知道是否有一個更簡潔的解決方案？

Answer 1

您可以使用，例如：

I = np.array([[1,2,3], [2,1,3], [4,1,6]]) 
M = np.ndarray((10,10,10,10)) 
X = np.array([M[t,:] for t in I]) 

Answer 2

這將是一個辦法做到這一點 -

import numpy as np 

# Get row indices for use when M is reshaped to a 2D array of d-columns format 
row_idx = np.sum(I*np.append(1,np.cumprod(M.shape[1:-1][::-1]))[::-1][:,None],0) 

# Reshape M to d-columns 2D array and use row_idx to get final output 
out = M.reshape(-1,M.shape[-1])[row_idx] 

作爲，另發現row_idx，如果你想避免np.append，你可以做 -

row_idx = np.sum(I[:-1]*np.cumprod(M.shape[1:-1][::-1])[::-1][:,None],0) + I[-1] 

或者少一點可怕獲取方式row_idx -

_,p2,p3,_ = M.shape 
row_idx = np.sum(I*np.array([p3*p2,p3,1])[:,None],0)