Объединение / суммирование строк numy массива на основе списка индексов

Question

У меня есть массив, в котором я хотел бы добавить определенные c строк вместе, чтобы получить массив с меньшим количеством строк.

import numpy as np
a = np.arange(50).reshape(10,5)
b = [0,0,0,1,1,2,2,2,2,3]
a
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24],
       [25, 26, 27, 28, 29],
       [30, 31, 32, 33, 34],
       [35, 36, 37, 38, 39],
       [40, 41, 42, 43, 44],
       [45, 46, 47, 48, 49]])

a.shape
(10, 5)

len(b)
10

Я хочу использовать b в качестве основы для ряды, которые я хочу объединить. Первые 3 строки будут суммированы, чтобы стать новой 1-й строкой, 4-я и 5-я строки будут суммированы, чтобы стать новой 2-й строкой. И так далее и тому подобное.

Желаемый результат:

array([[15, 18, 21, 24, 27],
       [35, 37, 39, 41, 43],
       [130, 134, 138, 142, 146],
       [45, 46, 47, 48, 49]])

Цикл неэффективен для моего кошелька. Я не уверен, можно ли это сделать в numpy, но, возможно, pandas или xarray?

Любая помощь приветствуется.

Answer 1

В pandas решением является создание DataFrame и совокупная сумма по индексу, созданному массивом b:

a = np.arange(50).reshape(10, 5)
b = [0,0,0,1,1,2,2,2,2,3]
print (a)

c = pd.DataFrame(a, index=b).sum(level=0).to_numpy()
print (c)
[[ 15  18  21  24  27]
 [ 35  37  39  41  43]
 [130 134 138 142 146]
 [ 45  46  47  48  49]]

Answer 2

Я хотел бы добавить решение numpy:

import numpy as np

a = np.arange(50).reshape(10,5)
b = [0,0,0,1,1,2,2,2,2,3]

sum_common = lambda x : sum(a[b==x,:])

indx = np.unique(b)

c = np.array(map(sum_common, indx))

Конечно, это можно сделать одной строкой:

c=np.array(map(lambda x : sum(a[b==x,:]), np.unique(b)))

Результат:

array([[ 15,  18,  21,  24,  27],
       [ 35,  37,  39,  41,  43],
       [130, 134, 138, 142, 146],
       [ 45,  46,  47,  48,  49]])