Fancy indexing of a numpy ndarray

user_8425270 0
Viewed 65

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 jth dimension.

How can i extractb from a ?

numpy python
2 Answers

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)

user_42384080

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]
user_32938810
Related