[BOJ] 2015 수들의 합

https://www.acmicpc.net/problem/2015

2015번: 수들의 합 4

---

문제

A[1], A[2], ..., A[N]의 N개의 정수가 저장되어 있는 배열이 있다. 이 배열 A의 부분합이란 1 ≤ i ≤ j ≤ N인 정수 i와 j에 대해 A[i]부터 A[j]까지의 합을 말한다.

N과 A[1], A[2], ..., A[N]이 주어졌을 때, 이러한 N×(N+1)/2개의 부분합 중 합이 K인 것이 몇 개나 있는지를 구하는 프로그램을 작성하시오.

---

입력

첫째 줄에 정수 N과 K가 주어진다. (1 ≤ N ≤ 200,000, |K| ≤ 2,000,000,000) N과 K 사이에는 빈칸이 하나 있다. 둘째 줄에는 배열 A를 이루는 N개의 정수가 빈 칸을 사이에 두고 A[1], A[2], ..., A[N]의 순서로 주어진다. 주어지는 정수의 절댓값은 10,000을 넘지 않는다.

---

출력

첫째 줄에 합이 K인 부분합의 개수를 출력한다.

---

풀이

1. input을 받으면서 누적 합을 저장한다.
2. N만큼 반복문을 돌면서 누적 합이 K인 경우 ans+1을 해준다. 
3. MAP[psum[i]-K]가 존재할 경우 그 만큼 ans에 더해준다.
4. MAP[psum[i]]+=1을 해준다.

여기서 MAP[psum[i]-K]가 의미하는 것은 psum[i]-psum[j]=K를 만족하는 j의 개수이다.

---

코드

#include <string>
#include <stdio.h>
#include <math.h>
#include <iostream>
#include <vector>
#include <algorithm>
#include <map>
#include <queue>
#include <stack>
#include <climits>
#include <set>

#define SWAP(a, b, type) do { \
    type temp; \
    temp = a;  \
    a = b;     \
    b = temp;  \
} while (0)

#define INF 1000000000
#define endl '\n'
#define ll long long

using namespace std;

ll N,K;
map<int, ll > MAP;
ll psum[200001];
ll ans=0;

void input() {
    cin>>N>>K;
    for(int i=1;i<=N;i++){
        cin>>psum[i];
        psum[i]+=psum[i-1];
    }
}

void init() {}

void solution() {
    for(int i=1;i<=N;i++){
        if(psum[i]==K) ans++;
        
        ans+=MAP[psum[i]-K];
        
        MAP[psum[i]]++;
    }
    
    cout<<ans;
}

int main()
{
    ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    input();
    solution();
    return 0;
}