Sum array by number in numpy

前端 未结 8 1157
别跟我提以往
别跟我提以往 2020-11-30 08:51

Assuming I have a numpy array like: [1,2,3,4,5,6] and another array: [0,0,1,2,2,1] I want to sum the items in the first array by group (the second array) and obtain n-groups

8条回答
  •  半阙折子戏
    2020-11-30 09:36

    The numpy function bincount was made exactly for this purpose and I'm sure it will be much faster than the other methods for all sizes of inputs:

    data = [1,2,3,4,5,6]
    ids  = [0,0,1,2,2,1]
    
    np.bincount(ids, weights=data) #returns [3,9,9] as a float64 array
    

    The i-th element of the output is the sum of all the data elements corresponding to "id" i.

    Hope that helps.

提交回复
热议问题