Numpy: Index 3d array by z-index array

Question

Say I have an array

a = np.zeros((3, 3, 3))

and a z-index array

z = np.random.randint(0, 3, (3, 3))

say z is

array([[1, 0, 2],
       [2, 2, 1],
       [1, 1, 0]])

Now I want to select the values of a with coordinates (starting at top left of z and traversing the array row-wise. Column-wise would be fine too.) (0, 0, 1), (0, 1, 0), (0, 2, 2), ... . The bold values are from the z array.

is a supposed to be a 3d array?

M4rtini
– M4rtini

2014-04-04 10:06:26 +00:00
Commented Apr 4, 2014 at 10:06 — M4rtini
– M4rtini, Commented Apr 4, 2014 at 10:06
Yes, sorry. Fixed it.

tauran
– tauran

2014-04-04 10:08:09 +00:00
Commented Apr 4, 2014 at 10:08 — tauran
– tauran, Commented Apr 4, 2014 at 10:08

M4rtini · Accepted Answer · 2014-04-04 10:32:53Z

1

I think this does what you want.

import numpy as np 

a = np.arange(3*3*3).reshape(3,3,3)
z = np.array([[1, 0, 2],
       [2, 2, 1],
       [1, 1, 0]])
i = np.arange(a.shape[0]).repeat(a.shape[0]).reshape(a.shape[0], a.shape[1])
j = i.T
#Should do the same as this. possible more efficient, did not test.
#i, j = np.indices(a[...,0].shape)
print a[i,j,z]

edited Apr 4, 2014 at 10:32

answered Apr 4, 2014 at 10:14

M4rtini

13.6k4 gold badges37 silver badges42 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

tauran Over a year ago

Nope. If I use zeros and then a[i,j,z] = 1, a[0, 0 ,0] is 1 but should be 0.

tauran Over a year ago

It works if i is array([[0, 0, 0], [1, 1, 1], [2, 2, 2]]) and j is i.T.

Collectives™ on Stack Overflow

Numpy: Index 3d array by z-index array

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related