1

I am trying to sum the different values of an array, a, that corresponds with the values of b.

a = np.array([2,3,4,5,6,7,8,9,10])
b = np.array([1,1,1,2,3,4,5,5,6])

For example, array b has 3 elements that are 1 (first 3 elements) and I want to add all the numbers in array a corresponding to this value. Also, 8 and 9 in array a should be added together since their b value is the same: 5, in this case. So the end result should be an array:

sum1 = np.array([9,5,6,7,17,10])

I have tried using a for loop with if-else statements, but I can't get the result that I'm looking for.

for i in range(1, len(a)-1):

if b[i-1] == b[i]:
    
    sum1[i-1] = np.sum(a[i-1:i])
    
else:
    sum1[i-1] = a[i-1]

So I was wondering if someone can help me with this by giving out an algorithm or pseudo-code that can help me. I have a text file with thousands of data and this is just a simplified version of what I want to do.

theheretic
  • 57
  • 6

1 Answers1

1

Try itertools.groupby:

from itertools import groupby

a = np.array([2, 3, 4, 5, 6, 7, 8, 9, 10])
b = np.array([1, 1, 1, 2, 3, 4, 5, 5, 6])

c = [sum(v for v, _ in g) for v, g in groupby(zip(a, b), lambda x: x[1])]
print(c)

Prints:

[9, 5, 6, 7, 17, 10]
Andrej Kesely
  • 168,389
  • 15
  • 48
  • 91