백준[2015] 수들의 합 4

문제 설명

문제 링크

http://icpc.me/2015

풀이

누적합을 이용하여 풀 수 있다.

0~i번째까지의 합을 미리 구해놓고, 구해놓은 누적합을 잘 관찰하여 풀 수 있다.

psum[i] : 0 ~ i까지의 합이라고 할 때, i ~ j까지의 합은 psum[i] - psum[j-1]이다.

이때 구간 합이 k가 되는지 확인할 수 있는 여부는, psum[i] - psum[j-1] (j < i) == k가 되는 j 값이 몇개 존재하는지 확인하면 된다.

그리고 psum[i] 자체가 k인지도 체크하면 된다.

그렇기 때문에 psum[1] ~ psum[n]까지 값을 기록해야 하는데 값의 범위가 20억으로 매우 크므로 그냥 배열을 만들 수 없어, map을 이용하여 저장한다.

코드

#include <cstdio>
#include <algorithm>
#include <map>

using namespace std;
using ll = long long;

ll psum[200005];
map<ll, ll> m;

int main() {
    ll n, k; scanf("%lld %lld", &n, &k);
    for (ll i = 1; i <= n; i++) {
        ll tmp; scanf("%lld", &tmp);
        psum[i] = psum[i - 1] + tmp;
    }
    ll ans = 0;
    for (ll i = 1; i <= n; i++) {
        if (psum[i] == k) ans++;
        ans += m[psum[i] - k];
        m[psum[i]]++;
    }
    printf("%lld", ans);
}