BOJ 2015 수들의 합 4

www.acmicpc.net/problem/2015

2015번: 수들의 합 4

각각 i부터 j까지의 부분합 중 k인 수가 몇 개인지 세는 문제이다.

일단 prefix sum을 통해서 배열을 채웠는데

	for (int i = 1; i <= n; i++) {
		cin >> x; d[i] += d[i - 1] + x;
	}

이 후가 문제이다. 여기서 각각의 원소에 대해서 \(O(N^2)\) 으로 짜면 시간 초과가 난다.

따라서 i 부터 j 사이의 부분합의 개수를 Log N으로 찾아야 한다.

여기서 map이 사용된다.

	for (int i = 1; i <= n; i++) {
		if (d[i] == k) cnt++;
		cnt += M[d[i] - k];
		M[d[i]]++;
	}

한 줄씩 살펴보면,

 

		if (d[i] == k) cnt++;

d[i]가 k와 같으면 카운트 증가.

		cnt += M[d[i] - k];

현재까지의 누적합에서 k를 뺀 값을 키값으로 해서 카운트를 증가시켜준다. 즉

d [i] - d [j] == k 인 값을 찾는 것이다.

#include <iostream>
#include <map>
using namespace std;

int d[200001];
map<int, long long> M;

int main() {
	ios::sync_with_stdio(false); cin.tie(0);

	long long n, k, x, cnt = 0;
	cin >> n >> k;
	for (int i = 1; i <= n; i++) {
		cin >> x; d[i] += d[i - 1] + x;
	}
	for (int i = 1; i <= n; i++) {
		if (d[i] == k) cnt++;
		cnt += M[d[i] - k];
		M[d[i]]++;
	}
	cout << cnt << endl;

	return 0;
}