CODEWITHSUNDEEP
pythonnumpy
Axel Puig
Axel Puig

Reputation: 1334

Reduce the size of a numpy array by splitting it and summing between the splits

I have an array like this:

1 2 3 4
2 3 4 5
3 4 5 6

I want to split it and sum the sub-arrays to get a smaller array.

For example, i want to split it given the following group sizes: [[2, 1], [1, 3]]

1 | 2 3 4
2 | 3 4 5
---------
3 | 4 5 6

The result should be:

1+2  2+3+4+3+4+5
3    4+5+6


3 21
3 15

Is there an efficient way to do this with numpy functions instead of for loops? I need to do it on high dimensional arrays. Thanks!

Upvotes: 3

Views: 258

Answers (1)

mathfux
mathfux

Reputation: 5949

Suppose you have a bigger array:

arr = np.random.randint(10, size=(6, 8))
print(arr)
[[3 5 1 9 7 3 0 1]
 [7 8 3 7 7 8 8 4]
 [3 0 2 2 2 5 4 9]
 [3 5 8 5 7 8 3 8]
 [3 2 7 0 9 4 0 8]
 [9 9 4 1 8 9 4 2]]

Then you can call np.add.reduceat two times on different axis:

def add_blocks(arr, s1, s2):
    s1, s2 = [0] + s1[:-1], [0] + s2[:-1]
    arr1 = np.add.reduceat(arr, np.cumsum(s1), axis=0)
    arr2 = np.add.reduceat(arr1, np.cumsum(s2), axis=1)
    return arr2

add_blocks(arr, [2, 3, 1], [1, 1, 1, 5])
array([[10, 13,  4, 54],
       [ 9,  7, 17, 74],
       [ 9,  9,  4, 24]], dtype=int32)

Upvotes: 6

Related Questions