NumPy Array Indexing

Tags:

numpy

Simple question here about indexing an array to get a subset of its values. Say I have a `recarray` which holds ages in one space, and corresponding values in another. I also have an array which is my desired subset of ages. Here is what I mean:

``````ages = np.arange(100)
values = np.random.uniform(low=0, high= 1, size = ages.shape)
data = np.core.rec.fromarrays([ages, values], names='ages,values')
desired_ages = np.array([1,4, 16, 29, 80])
``````

What I'm trying to do is something like this:

``````data.values[data.ages==desired_ages]
``````

But, it's not working.

570

mike

You want to create an subarray containing only the values whose indexes are in `desired_ages`.

Python doesn't have any syntax that directly corresponds to this, but list comprehensions can do a pretty nice job:

``````result = [value for index, value in enumerate(data.values) if index in desired_ages]
``````

However, doing it this way results in Python scanning through `desired_ages` for each element in `data.values`, which is slow. If you could insert

``````desired_ages = set(desired_ages)
``````

on the line before, this would improve performance. (You can determine if a value in is a set in constant time, regardless of the set's size.)

Complete Example

``````import numpy as np

ages = np.arange(100)
values = np.random.uniform(low=0, high= 1, size = ages.shape)
data = np.core.rec.fromarrays([ages, values], names='ages,values')
desired_ages = np.array([1,4, 16, 29, 80])

result = [value for index, value in enumerate(data.values) if index in desired_ages]
print result
``````
Output
``````[0.45852624094611272, 0.0099713014816563694, 0.26695859251958864, 0.10143425810157047, 0.93647796171383935]
``````
167

Jeremy

I changed your example a little, shuffle the order of ages:

``````import numpy as np
np.random.seed(0)
ages = np.arange(3,103)
np.random.shuffle(ages)
values = np.random.uniform(low=0, high= 1, size = ages.shape)
data = np.core.rec.fromarrays([ages, values], names='ages,values')
desired_ages = np.array([4, 16, 29, 80])
``````

If all the elements of desired_ages are in data.ages, you can sort data by age field first, and then use searchsorted() to find all the index quickly:

``````data.sort(order="ages") # sort by ages
print data.values[np.searchsorted(data.ages, desired_ages)]
``````

or you can use np.in1d the get a bool array and use it as index:

``````print data.values[np.in1d(data.ages, desired_ages)]
``````
28