NumPy Array Indexing

Question

I have a rec-array which holds ages in one space, and corresponding values in another. I also have an array which is my desired subset of ages. For example:

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])

I'm trying to do is something like this:

data.values[data.ages==desired_ages]

But, it's not working. Can anyone tell me the solution to this?

Answer 1

If you want to create a 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)

This might help!