Hi I have a python code as:
import numpy as np
a = np.array([[1,2,3],[4,5,6],[7,8,9]])
b = a[0,1]+a[2,2]-a[1,1]
>>> b
>>> 6
Is there any other faster way to add/subtract a list of specific elements of an array?
1 Answer
If you simply want to retrieve the values in an array from a list of indices, and sum them, you can do:
import numpy as np
a = np.array([[1,2,3],[4,5,6],[7,8,9]])
ind = [[0,1],[2,2],[1,1]]
values = a[zip(*ind)]
b = values.sum()
# b = 2+9+5 = 16
Note that I set indices ind as a list of indices pairs, and thus zip is necessary. Obviously this can be done in multiple ways. a[...] just needs to get a list or tuple (not numpy array) of indices grouped by axis, i.e. a[(0,2,1),(1,2,1)].
Now to make arbitrary addition or subtraction, there are multiple ways possible. For example:
import numpy as np
a = np.array([[1,2,3],[4,5,6],[7,8,9]])
ind = [[0,1],[2,2],[1,1]] # indices in array
op = [1,1,-1] # 1 for addition, -1 for subtraction
values = a[zip(*ind)]*op
b = values.sum()
# b = 2+9-5 = 6
One last point: this method is useful for a set of indices of arbitrary size (i.e. which would be a parameter of your code). For 3 specific values, its is better to do it explicitly as in your code.
Comments
numpy.sumfor fast masking and addition on multi-dimensional arrays.