Vectorizing this for-loop in numpy

Question

I was wondering how I would vectorize this for loop. Given a 2x2x2 array x and an array where each element is the ith, jth, and kth element of the array I want to get x[i,j,k]

Given an arrays x and y

x = np.arange(8).reshape((2, 2, 2))
y = [[0, 1, 1], [1, 1, 0]]

I want to get:

x[0, 1, 1] = 3 and x[1, 1, 0] = 6

I tried:

print(x[y])

But it prints:

array([[2, 3],
       [6, 7],
       [4, 5]])

So I ended up doing:

for y_ in y:
    print(x[y_[0], y_[1], y_[2]])

Which works, but I can't help but think there is a better way.

akuiper · Accepted Answer · 2018-10-08 00:55:35Z

1

Use transposed y i.e zip(*y) as the index; You need to have the indices for each dimension as an element for advanced indexing to work:

x[tuple(zip(*y))]
# array([3, 6])

answered Oct 8, 2018 at 0:55

akuiper

216k33 gold badges362 silver badges379 bronze badges

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

Collectives™ on Stack Overflow

Vectorizing this for-loop in numpy

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related