[백준/C++] 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인 부분합의 개수를 출력한다.

풀이

01. 일단 자료형에 유의하자.

- 부분합 개수의 최댓값이 2,000,000,000 * (2,000,000,000 + 1) / 2 이므로 long long을 사용해야 한다.

- 결과값의 최댓값은 부분합 개수의 최댓값과 같으므로 역시 long long을 사용한다.

 

02. 각각의 부분합을 저장한다(p_sum)

 

03. 1~n까지 i를 돌면서 sum[i] - sum[j] = k를 만족하는 j가 있는지 검사한다. 

문제에서는 1<i<j<n으로 되어있지만 풀이 과정에서는 i랑 j를 구분할 필요가 없기 때문에 (1<j<i<n)으로 바꿔썼다.

- sum[j] = sum[i] - k

 

04. sum[i]값을 map에 넣어준다. 

- map은 부분합의 빈도를 저장한다. (key : 부분합, value : 빈도)

- 만약 sum[j]가 map에 존재하면 ans를 증가시켜준다.

- p_sum[i]가 k인 경우, j가 0이 될 수 있기 때문에 map[0] = 1을 표시해준다.

 

#include <iostream>
#include <map>
using namespace std;
int p_sum[200001];

int main(){
    int n, k, num;
    cin >> n >> k;
    map<int, long long> m; // 부분합 개수 : 2000000000*(2000000000+1)/2

    long long cnt = 0;
    m[0] = 1; // p_sum[i]가 k인 경우, j=0이 될 수 있음

    for (int i=1; i<=n; i++){
        cin >> num;
        // 01. 각각의 부분합을 저장
        p_sum[i] = p_sum[i-1] + num;
        // * 1~n까지 i를 돌면서 sum[i] - sum[j] = k를 만족하는 j가 있는지 검사(i>=j라고 가정! - 문제와 기호 다르게 써도 상관X)
        // 만약 존재하면 ans++;
        // sum[j] = sum[i] - k
        cnt += m[p_sum[i] - k];
        // 02. map에 저장
        m[p_sum[i]]++;
    }
    cout << cnt;
}

---

이 문제는 이해를 못해서 너무 어려웠다. 특히 다 풀었는데 p_sum[i]가 k인 경우를 고려하지 못해서 자꾸 에러가 났다.