I have a 3d array a = np.arange(108).reshape(6, 6, 3). I want to grab certain indices of the array, as defined by i = np.array([[0, 1], [1, 3], [2, 1]]) such that the result is [[3, 4, 5], [27, 28, 29], [39, 40, 41]]. I need an efficient way to do this, as my actual arrays are significantly larger.
Add a comment
|
1 Answer
Extract the first and second dimension indices from i, then use advanced indexing:
a[i[:,0], i[:,1], :] # or a[i[:,0], i[:,1]]
#array([[ 3, 4, 5],
# [27, 28, 29],
# [39, 40, 41]])
