Fancy indexing of a numpy ndarray

Question

Suppose i have an array shaped as a :

import numpy as np
n = 10
d = 5
a = np.zeros(shape = np.repeat(n,d))

And that I want to obtain the values corresponding to indexes (0,...,:,...,0) for the : along dimensions, resulting in a (n,d) -shaped array b , with b[i,j] = a[0,...,0,i,0,...,0] where the i is in the j th dimension.

How can i extract b from a ?

Answer 1

Easiest is to do a for loop:

# get the first slice of `a` along given dimension `j`
def get_slice(a,j):
    idx = [0]*len(a.shape)
    idx[j] = slice(None)
    return a[tuple(idx)]

out = np.stack([get_slice(a,j) for j in range(len(a.shape))])

And out.shape is (10,5)

Answer 2

Get the flattened indices and just index for a vectorized solution -

n = len(a)
d = a.ndim
idxs = np.multiply.outer(n**np.arange(d), np.arange(n))
out = a.flat[idxs]