Сумма строк двумерного массива np со списком индексов

Question

У меня есть массив 2d numpy и список чисел. Если список равен [1, 3, 1, 8], где список суммируется с количеством строк, я хочу вывести массив с неизменной первой строкой, суммированием следующих трех строк, неизменением пятой строки и суммированием оставшихся восьми строк.

В качестве примера:

A = [[0,0], [1,2], [3,4]] и l = [1, 2] будут выводить [[0,0], [4,6]

Я просмотрел np.sum и другие функции, но не смог найти не эта функциональность. Спасибо.

Answer 1

Если количество элементов в l относительно велико, вы можете повысить производительность, используя groupby из pandas, например,

import pandas as pd

labels = np.repeat(np.arange(1, len(l) + 1), l)
# [1, 2, 2]
df = pd.DataFrame(A)
df['label'] = labels
result = df.groupby('label').sum().values

Answer 2

Я закончил тем, что придумал свое собственное решение, когда понял, что могу отсортировать свой список, не влияя на желаемый результат. Я использовал np.unique, чтобы определить первые индексы каждого элемента в отсортированном списке, а затем суммировал строки между этими индексами. Смотри ниже.

elements, indices = np.unique(data, return_counts=True)
row_summing = np.append([0], np.cumsum(indices))[:-1] #[0, index1, index2,...]

output = np.add.reduceat(matrix, row_summing, axis=0)

Answer 3

Вы можете просто выполнить итерации по индексам l и на основе позиции либо взять эту строку, либо сумму по диапазону строк.

import numpy as np

A = [[0,0], [1,2], [3,4]]
l = [1, 2]

ans = []
for i in range(len(l)):
  if i%2 == 0:
    ans.append(A[ l[i] ])
  else:
    ans.append( np.sum( A[ l[i-1]:l[i-1] + l[i] ], axis=0 ) )

ans = np.array(ans)
print(ans)

[[1 2]
 [4 6]]

Примечание:

Если список [1, 3, 1, 8], где список суммируется с количеством строк, я хочу вывести массив с первой строкой без изменений, с суммированием следующих трех строк, без изменения пятой строки, и оставшиеся восемь строк суммируются.

• Я думаю, вы имели в виду [1, 3, 5, 8]