Суммы фрагментов переменного размера в списке, где размеры заданы другим списком - PullRequest
8 голосов
/ 07 мая 2020

Я хотел бы получить следующую сумму из двух списков:

a = [0,1,2,3,4,5,6,7,8,9]
b = [2,3,5]

Результат должен быть суммой каждого b элемента a, например:

  1. b[0] = 2, поэтому результат первой суммы должен быть: sum(a[0:2])
  2. b[1] = 3, поэтому результат второй суммы должен быть: sum(a[2:5])
  3. b[2] = 5, поэтому результат третьей суммы должно быть: sum(a[5:10])

Результат печати: 1,9,35

Ответы [ 5 ]

9 голосов
/ 07 мая 2020

Вы можете использовать np.bincount с weights:

groups = np.repeat(np.arange(len(b)), b)

np.bincount(groups, weights=a)

Вывод:

array([ 1.,  9., 35.])
6 голосов
/ 07 мая 2020

NumPy имеет инструмент для сокращения суммы на основе срезов с помощью np.add.reduceat -

In [46]: np.add.reduceat(a,np.cumsum(np.r_[0,b[:-1]]))                          
Out[46]: array([ 1,  9, 35])
3 голосов
/ 07 мая 2020

Трудно конкурировать с решением np.bincount, но вот еще один хороший способ приблизиться к нему с помощью np.cumsum:

strides = [0] + np.cumsum(b).tolist()  # [0, 2, 5, 10]
stride_slices = zip(strides[:-1], strides[1:])  # [(0, 2), (2, 5), (5, 10)]
[sum(a[s[0]: s[1]]) for s in stride_slices]
# [1, 9, 35]
3 голосов
/ 07 мая 2020

Вы имеете в виду что-то вроде этого?

a = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
b = [2, 3, 5]

def some_function(a, b): # couldnt come up with a name :D
    last_index = 0

    for i in b:
        print(sum(a[last_index:last_index + i]))
        last_index += i

some_function(a, b)
0 голосов
/ 07 мая 2020

Вы можете использовать понимание списка с sum:

a=[0,1,2,3,4,5,6,7,8,9]
b=[2,3,5]
r = [sum(a[(k:=sum(b[:i])):k+j]) for i, j in enumerate(b)]

Вывод:

[1, 9, 35]
Добро пожаловать на сайт PullRequest, где вы можете задавать вопросы и получать ответы от других членов сообщества.
...