python numpy `np.take` with 2 dimensional array

Question

I'm trying to take a list of elements from an 2D numpy array with given list of coordinates and I want to avoid using loop. I saw that np.take works with 1D array but I can't make it work with 2D arrays.

Example:

a = np.array([[1,2,3], [4,5,6]])
print(a)
# [[1 2 3]
#  [4 5 6]]

np.take(a, [[1,2]])
# gives [2, 3] but I want just [6]

I want to avoid loop because I think that will be slower (I need speed). But if you can persuade me that a loop is as fast as an existing numpy function solution, then I can go for it.

Does this do what you want? a[1].take([2]) Ignoring the index error in your question... this returns array([6]) — kpie
– kpie, Commented Feb 12, 2020 at 15:47
Then, if I have a list of coordinates, I need to do [a[i].take(j) for i,j in coords], will that be slower if there's a built-in numpy function that can potentially do it? — Yilun Zhang
– Yilun Zhang, Commented Feb 12, 2020 at 15:51
How your list of coordinates looks like? Is it like ‘[[y0, x0], [y1, x1], ...]‘ ? — kuzand
– kuzand, Commented Feb 12, 2020 at 16:03

kuzand · Accepted Answer · 2020-02-12 16:35:36Z

3

If I understand it correctly, you have a list of coordinates like this:

coords = [[y0, x0], [y1, x1], ...]

To get the values of array a at these coordinates you need:

a[[y0, y1, ...], [x0, x1, ...]]

So a[coords] will not work. One way to do it is:

Y = [c[0] for c in coords]
X = [c[1] for c in coords]

or

Y = np.transpose(coords)[0]
X = np.transpose(coords)[1]

Then

a[Y, X]

answered Feb 12, 2020 at 16:35

kuzand

9,8864 gold badges48 silver badges50 bronze badges

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

1 Comment

Yilun Zhang Over a year ago

Thank you for the answer! Exactly what I wanted. I just tested out and it seems like the speed is very similar to np.take method along with list comprehension. But still kudos for you!

Tls Chris · Accepted Answer · 2020-02-12 16:37:54Z

1

Does fancy indexing do what you want? np.take seems to flatten the array before operating.

import numpy as np

a = np.arange(1, 10).reshape(3,3)

a
# array([[1, 2, 3],
#        [4, 5, 6],
#        [7, 8, 9]])

rows = [ 1,1,2,0]
cols = [ 0,1,1,2]

# Use the indices to access items in a
a[rows, cols]
# array([4, 5, 8, 3])

a[1,0], a[1,1], a[2,1], a[0,2]
# (4, 5, 8, 3)

answered Feb 12, 2020 at 16:37

Tls Chris

4,0461 gold badge12 silver badges28 bronze badges

3 Comments

kuzand Over a year ago

How is this different from what I wrote? ;)

Tls Chris Over a year ago

It was written in parallel. I only saw yours after I posted.

Yilun Zhang Over a year ago

Thank you for the answer! Exactly what I wanted. I just tested out and it seems like the speed is very similar to np.take method along with list comprehension. Since Andreas K. answered slight faster than you, so I will accept his answer, but still kudos for you!

Collectives™ on Stack Overflow

python numpy `np.take` with 2 dimensional array

2 Answers 2

1 Comment

3 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

3 Comments

Your Answer

Sign up or log in

Post as a guest

Related