[BOJ] 수들의 합 4 2015번 C++

 

문제

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인 부분합의 개수를 출력한다.

 

접근 방식

N의 최대값이 20만이기 때문에

누적합을 사용해도

i,j를 구하는 2중 반복문을 사용하면 시간초과가 발생한다..!

 

그래서 찾아보다가 map을 사용해 구현하는 방식에 대해 알게 되었다.

일단 누적합 배열 pSum을 만들어 i까지의 누적합을 저장한다.

counts 해시 테이블은 누적 합이 몇 번 나타났는지를 저장하는 자료구조다

counts[pSum[i]-K]를 확인해서 만약 값이 존재하면 ans에 저장해두는데

여기서 counts[pSum[i]-K]는 현재 지점까지의 누적 합에서 K를 뺀값이 이전에 존재하는지를 확인하는 것이다. 

 

만약 

A = [3, 1, -4, 2, 5], K = 3 일때

pSum은 [3, 4, 0, 2, 7]이고 

1. i=1: pSum[1] = 3, counts[3 - 3]을 조회, counts[0] = 1이므로 ans = 1이 됨
2. i=2: pSum[2] = 4, counts[4 - 3]을 조회, counts[1]은 존재하지 않음. 따라서, 변화 없음.
3. i=3: pSum[3] = 0, counts[0 - 3]은 존재하지 않음 . 따라서, 변화 없음.
4. i=4: pSum[4] = 2, counts[2 - 3]은 존재하지 않음 . 따라서, 변화 없음.
5. i=5: pSum[5] = 7, counts[7 - 3]을 조회합니다. counts[4]는 존재하지 않음, 변화 없음.

이러한 결과로 K가 3일때 결과값은 1이된다

 

 

누적합 정말 어려운 것 같다...

 

코드

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
int N,K;
int arr[200002];
ll pSum[2000002];
ll ans = 0;
unordered_map<ll, int> counts;

int main() {
    ios::sync_with_stdio(false); cin.tie(0);
    cin>>N>>K;
    counts[0] = 1;
    for(int i=1; i<=N; i++) {
        cin>>arr[i];
        pSum[i] = pSum[i-1] + arr[i];
        ans += counts[pSum[i]-K];
        counts[pSum[i]]++;
    }

    cout<<ans;

}