累積和の練習でAOJ Maximum Sum を解く

累積和の練習で、Aizu Online JudgeのMaximum Sumを解いて解説します。

https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=0516

n個の整数からなる数列と整数kが与えられ、
連続して並ぶk個の整数の和の最大値を求める問題です。

nはで、kはです。
単純に頭からk個の整数を読み取りながら足していくと、
n - k 箇所でk回の足し算が必要になり、計算量はになり、アウトです。
そこで、累積和の考え方でk回の足し算をしないように工夫します。

accを累積和を格納する配列とした時に、acc[i]をインデックス0からiまでの累積和になるようにします。

acc[0] = a_arr[0]
for i in range(1, n):
    acc[i] = acc[i - 1] + a_arr[i]

この累積和の配列を使って、インデックスiから i + kまでの和は以下の通りです。

acc[i + k] - acc[i]

連続して並ぶk個の整数の和の最大値を求めたいので、以下のようにします。

ans = max(ans, acc[j + k] - acc[j])

最後のコードは以下の通りです。

def solve(n: int, k: int, a_arr: list) -> int:
    acc = [0 for _ in range(n)]
    acc[0] = a_arr[0]
    for i in range(1, n):
        acc[i] = acc[i - 1] + a_arr[i]

    j = 0
    ans = 0
    while j + k < n:
        ans = max(ans, acc[j + k] - acc[j])
        j += 1
    return ans


while True:
    n, k = map(int, input().split())
    if n == 0 and k == 0:
        exit()
    a_arr = [int(input()) for _ in range(n)]
    print(solve(n, k, a_arr))

2日連続で簡単な累積和の問題を解きました。
次はもう少し難しめの累積和の問題にトライしようと思います。