Посчитать префиксные суммы с конца. Python
Решаю задачу. Есть массив префиксных сумм [0,2,5,10,13] (для чисел 2,3,5,3 соответственно). Мне нужно посчитать сумму последних k элементов. Допустим k = 2. Значит с конца массива нужно посчитать (13 - 10) + (10-5) = 8. Не могу реализовать это в коде, как сделать? Вот мой код:
n = int(input())
a = []
rez = []#ответы
pref = [0]#суммы
for i in range(n):
b = input()
if b[0] == "-":#если просят удалить
rez.append(a.pop())
elif b[0] == "+":#если просят добавить
a.append(int(b[1:]))
pref.append(pref[-1]+int(b[1:]))
elif b[0] == "?":#если просят вывести сумму последних k эдементов
k = int(b[1:])
#print(a)
rez.append(sum(a[-k:]))
for i in rez:#выводим ответы
print(i)
Ответы (2 шт):
Автор решения: strawdog
→ Ссылка
Вам наверняка предложат "ванильное" решение, но мне это не интересно, поэтому вот вам с numpy, например:
import numpy as np
sums = [0,2,5,10,13] # ваши суммы
k = 2 # ваше k
res = sum(np.diff(sums,1)[-k:])
print(res)
# 8
Автор решения: CrazyElf
→ Ссылка
Ну вот пусть будет без Numpy, на основании кода @strawdog и @MaxU:
sums = [0,2,5,10,13] # ваши суммы
k = 2 # ваше k
res = sum([y - x for x, y in zip(sums[-k-1:], sums[-k:])])
print(res)
# 8
Хотя да, прав @extrn, на самом деле суммы тут вообще не нужны, нужно просто два элемента вычесть, только сообразить какие.