Посчитать префиксные суммы с конца. Python

Question

Решаю задачу.
Есть массив префиксных сумм [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)

CrazyElf · Answer

Ну вот пусть будет без 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, на самом деле суммы тут вообще не нужны, нужно просто два элемента вычесть, только сообразить какие.

strawdog · Answer

Вам наверняка предложат "ванильное" решение, но мне это не интересно, поэтому вот вам с numpy, например:
import numpy as np

sums = [0,2,5,10,13] # ваши суммы
k = 2 # ваше k

res = sum(np.diff(sums,1)[-k:])
print(res)
# 8

Посчитать префиксные суммы с конца. Python

2 Answers

Add your own answers!

Ask a Question