fen
fen

Reputation: 91

Sums of arrays in an array

Is there a speedy way to do this;

import numpy as np
a=np.array([1,2,3,4])
b=np.array([1,2])
c=np.array([a,b])
result=magic(c)

where magic() is the functionality I want and result should be np.array([10,3]) i.e. a numpy.array containing the sums of each of the input arrays.

Upvotes: 2

Views: 99

Answers (2)

askewchan
askewchan

Reputation: 46530

Here's a compilation of the suggestions (answers and comments) and their timings:

import numpy as np

c = np.array([np.random.rand(np.random.randint(1, 300)) for i in range(50)])

def oliver(arr):
    res = np.empty_like(arr)
    for enu, subarr in enumerate(arr):
        res[enu] = np.sum(subarr)
    return res

def reut(arr):
    return np.array([a.sum() for a in arr])

def hpaulj(arr):
    d = np.concatenate(arr)
    l = map(len, arr)
    i = np.cumsum(l) - l
    return np.add.reduceat(d, i)

And their times:

In [94]: timeit oliver(c)
1000 loops, best of 3: 457 µs per loop

In [95]: timeit reut(c)
1000 loops, best of 3: 317 µs per loop

In [96]: timeit hpaulj(c)
10000 loops, best of 3: 94.4 µs per loop

It was somewhat tricky to implement @hpaulj's, but I think I got it (and it's the fastest if you use concatenate instead of hstack)

Upvotes: 2

Oliver W.
Oliver W.

Reputation: 13459

Of the many possible solutions, this is one:

def your_magic(arr):
    res = np.empty_like(arr)
    for enu, subarr in enumerate(arr):
        res[enu] = np.sum(subarr)
    return res

Be mindful though that making a numpy array of unequal length arrays is not efficient at all and is just very similar to adding arrays in a normal Python list. This is the reason the returned array res in the function above will in general be of the dtype object.

Upvotes: 1

Related Questions