[백준 2015번] 수들의 합4

문제 링크:https://www.acmicpc.net/problem/2015 

2015번: 수들의 합 4

<문제 풀이> 누적합, 해시맵

 

이 문제는 누적합 배열을 만들고 각 i에 대해서 j < i 일 때 prefix[i] - prefix[j] = k를 만족하는 부분합의 개수를 찾는 문제입니다.

 

하지만 모든 i에 대해 j<i인 j에 대해 부분합을 하나씩 계산하면 O(N^2)의 시간복잡도를 가지기 때문에 시간 초과가 발생할 수 있습니다.

 

대신 prefix[i] - prefix[j] = k에서 우리는 prefix[i]와 k의 값을 이미 알고 있으므로

 

prefix[j] = prefix[i] - k로 바꿀 수 있습니다. 따라서, prefix[i] - k가 prefix[j]의 값과 일치하는 부분합의 개수를 찾는 문제로 변경됩니다.

 

해시맵을 이용하면 O(lgN)의 시간복잡도로 prefix[j]의 값과 일치하는 부분합의 개수를 알 수 있습니다. 이를 이용하여 전체 시작복잡도는 O(N*lgN)이 됩니다.

 

여기서 주의할 점 한 가지는 부분합의 개수가 N*(N+1)/2개이므로 최대 200000*(200000+1)/2 개의 부분합을 카운트해야 허가 때문에 int가 아닌 long long 자료형에 저장을 해야 합니다.

 

 

<C++소스코드>

 

#include <iostream>
#include <unordered_map>
using namespace std;
 
typedef long long ll;
 
int n, k;
int prefix[200001];
int arr[200001];
ll ans = 0;
unordered_map<int, ll> m;
int main() {
    cin >> n >> k;
    for (int i = 1; i <= n; i++) {
        cin >> arr[i];
        prefix[i] = prefix[i - 1] + arr[i];
    }
    m[0] = 1; // prefix[i] == k 
    for (int i = 1; i <= n; i++) {
        ans += m[prefix[i] - k];
        m[prefix[i]]++;
    }
    cout << ans;
 
}