FranGoitia FranGoitia - 22 days ago 7
Python Question

numpy. Sum multidimensional array at specific indexes

I'm trying to replace something like this code, for a vectorized efficient operation using numpy.

counter = 0
idxs = [1, 3]
lists = [[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]]
for l in lists:
for idx in idxs:
counter += l[idx]

Answer

Just sum the array:

idxs = [1, 3]
lists = [[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]]
ary = np.array(lists)
counter = ary[:,idxs].sum()