CODEWITHSUNDEEP
numpynumpy-slicing
michmest
michmest

Reputation: 25

Numpy slicing with list

I'm trying to slice an Ndarray a with a list b. But the behaviour is not as I would expect it. What do I have to change to get the wanted result?

a = np.arange(27).reshape(3,3,3)

b = [2,2]

Actual behaviour:

a[:,b]

array([[[ 6,  7,  8],
        [ 6,  7,  8]],

       [[15, 16, 17],
        [15, 16, 17]],

       [[24, 25, 26],
        [24, 25, 26]]])

Wanted behaviour:

a[:,2,2]

array([ 8, 17, 26])

Upvotes: 1

Views: 59

Answers (1)

Henry Ecker
Henry Ecker

Reputation: 35626

Try replacing : with slice(None) and unpacking b:

a[(slice(None), *b)]
# [ 8 17 26]

import numpy as np

a = np.arange(27).reshape((3, 3, 3))
b = [2, 2]
result = a[(slice(None), *b)]

print(result)
# [ 8 17 26]
print((result == a[:, 2, 2]).all())
# True

Upvotes: 3

Related Questions