CODEWITHSUNDEEP
pythonnumpy
user9413641
user9413641

Reputation:

numpy sum of each array in a list of arrays of different size

Given a list of numpy arrays, each of different length, as that obtained by doing lst = np.array_split(arr, indices), how do I get the sum of every array in the list? (I know how to do it using list-comprehension but I was hoping there was a pure-numpy way to do it).

I thought that this would work:

np.apply_along_axis(lambda arr: arr.sum(), axis=0, arr=lst)

But it doesn't, instead it gives me this error which I don't understand:

ValueError: operands could not be broadcast together with shapes (0,) (12,)

NB: It's an array of sympy objects.

Upvotes: 0

Views: 1046

Answers (1)

Péter Leéh
Péter Leéh

Reputation: 2119

There's a faster way which avoids np.split, and utilizes np.reduceat. We create an ascending array of indices where you want to sum elements with np.append([0], np.cumsum(indices)[:-1]). For proper indexing we need to put a zero in front (and discard the last element, if it covers the full range of the original array.. otherwise just delete the [:-1] indexing). Then we use the np.add ufunc with np.reduceat:

import numpy as np
arr = np.arange(1, 11)
indices = np.array([2, 4, 4])

# this should split like this
# [1  2 | 3  4  5  6 | 7  8  9  10]

np.add.reduceat(arr, np.append([0], np.cumsum(indices)[:-1]))
# array([ 3, 18, 34])

Upvotes: 2

Related Questions