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!
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)