ABC037 C - 総和
備忘録
問題
回答
import sys, os, math, bisect, itertools, collections, heapq, queue # from scipy.sparse.csgraph import csgraph_from_dense, floyd_warshall from decimal import Decimal from collections import defaultdict # import fractions sys.setrecursionlimit(10000000) ii = lambda: int(sys.stdin.buffer.readline().rstrip()) il = lambda: list(map(int, sys.stdin.buffer.readline().split())) fl = lambda: list(map(float, sys.stdin.buffer.readline().split())) iln = lambda n: [int(sys.stdin.buffer.readline().rstrip()) for _ in range(n)] iss = lambda: sys.stdin.buffer.readline().decode().rstrip() sl = lambda: list(map(str, sys.stdin.buffer.readline().decode().split())) isn = lambda n: [sys.stdin.buffer.readline().decode().rstrip() for _ in range(n)] lcm = lambda x, y: (x * y) // math.gcd(x, y) # lcm = lambda x, y: (x * y) // fractions.gcd(x, y) MOD = 10 ** 9 + 7 MAX = float('inf') def main(): if os.getenv("LOCAL"): sys.stdin = open("input.txt", "r") N, K = il() A = il() S = [0] * (N + 1) for n in range(N): S[n + 1] = A[n] + S[n] ret = 0 for n in range(K, N + 1): ret += S[n] - S[n - K] print(ret) if __name__ == '__main__': main()
考え方
予め、累積和を用いてi
番目までの合計を求めておき、
求めた合計から区間(連続する部分)の総和を求める。
与えられた数列a
の累積和(S
)を求める。
累積和を求めることによって、i
番目までの総和がわかるため、
区間(i, j)
の総和を求める場合にはS[j] - S[i]
で求めることができる。
今回の問題は長さK
の連続部分の総和を求めるため、
K
からN
までの間で、それぞれS[n] - S[n-K]
の総和を求めればよい。