Official

D - K項足し算 / K Integers Summations Editorial by tatyam


\(\displaystyle B_x = \sum_{i=x}^{x+K-1}A_i\) とします。
\(B_1, B_2, \dots, B_{N-K+1}\) をそれぞれ求めると、 \((N-K+1)K\) 回程度の足し算が必要になり、TLE してしまいます。
そこで、\(A\) の累積和 \(\displaystyle S_x = \sum_{i=1}^xA_i = S_{x-1} + A_x\) を前計算しておくと、\(B_x = S_{x+K-1}-S_{x-1}\)\(1\) 回の引き算で求められるようになり、時間計算量 \(\Theta(N)\) でこの問題を解くことができます。
また、\(B_{i+1}-B_{i}=A_{i+K}-A_i\) を利用しても \(\Theta(N)\) で解くことができます。

回答例 (C++)

#include <iostream>
#include <vector>
#include <numeric>
using namespace std;
using ll = long long;

int main(){
    ll N, K;
    cin >> N >> K;
    vector<ll> A(N);
    for(ll& a : A) cin >> a;
    A.push_back(0);
    exclusive_scan(A.begin(), A.end(), A.begin(), 0LL);
    for(ll i = K; i <= N; i++) cout << A[i] - A[i - K] << '\n';
}

回答例 (Python)

N, K = map(int, input().split())
A = list(map(int, input().split()))
A.insert(0, 0)
for i in range(N):
    A[i + 1] += A[i]
for i in range(K, N + 1):
    print(A[i] - A[i - K])

posted:
last update: