CODEWITHSUNDEEP
pythonnumpylist-comprehension
lrnv
lrnv

Reputation: 1106

Subsetting an array with another array of indexes

Suppose i have a (list | np.array) a of size (N,) and a (list | np.array) l of integers between 0 and N-1.

Is there a way i could write more efficiently sum([a[x] for x in l]) ?

Four different conditions:

Upvotes: 0

Views: 46

Answers (1)

alani
alani

Reputation: 13079

  • a is a numpy array, l is a numpy array
  • a is a numpy array, l is a list

For both of the above you can do a[l].sum()

  • a is a list, l is a numpy array
  • a is a list, l is a list

For these last two, your options are either to cast a to numpy and then do as above:

np.asarray(a)[l].sum()

or if you are going to use something like your list comprehension, then at least use a generator expression instead - there is no need to build a list simply to add up the values:

sum(a[x] for x in l)

If you are looking for a single solution that you can use regardless of the type, then np.asarray(a)[l].sum() (as suggested above) will work, because if the argument to np.asarray is an array anyway, then it will simply use it as-is -- but be aware that if a is a list then this will need to create an array version of a, so use of the generator expression will be more economical on memory.

import numpy as np

a_list = [10, 11, 12]
l_list = [2, 2]

a_array = np.array(a_list)
l_array = np.array(l_list)

for a in a_list, a_array:
    for l in l_list, l_array:
        print(np.asarray(a)[l].sum())

gives:

24
24
24
24

Upvotes: 1

Related Questions