CODEWITHSUNDEEP
pythonnumpy
Elhanan Schwarts
Elhanan Schwarts

Reputation: 383

Selecting multiple values from 3d numpy in efficient way

I have a very large 3d numpy from which I want to extract many values (x, y, z).

For the sake of simplicity let's say this is the numpy:

import numpy as np
a = np.arange(64).reshape(4,4,4)

From which I want to extract the values of the following collection of points:

points = [[3,0,0],[0,1,0],[3,0,1],[2,3,1]]

In this example, the expected result should be:

[48,4,49,45]

Because performance metter, I want to avoid iterate like the following code:

points  = [[3,0,0],[0,1,0],[3,0,1],[2,3,1]]
for i in points:
    print(a[i[0],i[1],i[2]])

Upvotes: 2

Views: 108

Answers (1)

swag2198
swag2198

Reputation: 2696

Try this. Uses numpy fancy/advanced indexing.

>>> import numpy as np
>>> a = np.arange(64).reshape(4,4,4)

>>> points = [[3,0,0],[0,1,0],[3,0,1],[2,3,1]]
>>> points = np.array(points)

>>> i = points[:, 0]
>>> j = points[:, 1]
>>> k = points[:, 2]
>>> a[i, j, k]
array([48,  4, 49, 45])

Upvotes: 3

Related Questions