문제

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

풀이

시간초과 때문에 누적합과 맵을 같이 사용해야 하는 문제였다. (+ long long 사용)

부분합이 K가 되는 경우를 세야 하는데 누적합을 이용하여 부분합을 구하기 때문에
누적합 차가 K가 되는 경우를 세야 한다.
하지만 누적합 배열을 돌면서 모든 경우를 세면 시간초과가 난다.

따라서 누적합의 원소들을 맵에 저장하면서 이전에 저장한 누적합 중에 차가 K가 되는 경우가 있었다면 ans에 더해준다.

처음에는 map에 누적합 원소들을 모두 저장해두고 차가 K가 되는 경우를 구하려고 했는데
이 경우에는 map 크기의 제곱 만큼 돌아야 하기 때문에 시간초과가 난다. (O(n^2))

따라서 map에 누적합 원소들을 저장하면서 동시에 이전에 저장한 차가 K가 되는 원소가 있었다면 ans에 더해주어야 한다. (O(n))

#include <iostream>
#include <algorithm>
#include <vector>
#include <string>
#include <map>

using namespace std;

int arr[200002];
int pSum[200002];
long long ans;

map<int, int> pSum_cnt;

int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);

    int N, K;
    cin >> N >> K;

    for(int i=1;i<=N;i++){
        cin >> arr[i];
        pSum[i] = pSum[i-1] + arr[i];
    }

    for(int i=1;i<=N;i++){
        if(pSum[i] == K) ans++;
        if(pSum_cnt[pSum[i]-K]) ans+=pSum_cnt[pSum[i]-K];
        pSum_cnt[pSum[i]]++;
    }

    cout << ans << "\n";

    return 0;
}