使用另一个3D索引数组索引4D数组

时间:2015-06-23 13:14:40

标签: numpy matrix-indexing

A有一个4D数组Ma x b x c x d)和一个数组I的索引(3 x f),例如

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]]但是,我想知道是否有更简洁的解决方案?

2 个答案:

答案 0 :(得分:1)

您可以使用use,例如:

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

答案 1 :(得分:0)

这是一种方法 -

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)