백준 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인 것이 몇 개나 있는지를 구하는 프로그램을 작성하시오.

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

---

 

 

 

 

 

 

제한 사항

---

 

 

 

 

 

 

풀이

문제를 요약하면 N개의 수가 주어질 때 i~j 번째의 수를 더한 값이 K가 되는 경우의 수를 구하면 된다

 

해당 문제는 완전 탐색으로 푼다면 굉장히 쉬운 문제이다.

하지만 완전 탐색으로 풀게 되면 시간 초과가 발생한다.

따라서 효율적으로 경우의 수를 구하는 방법이 필요하다.

 

우선 구간 합에 대한 질의이기 때문에 누적합을 사용하면 합을 구하는 시간을 줄일 수 있다.

하지만 누적합으로 구간 합을 구한다고 하더라고 i, j 쌍을 모두 구하게 되면 $O(N^2)$의 시간이 걸려 시간 초과가 발생한다.

따라서 더 빠르게 경우의 수를 구하는 방법이 필요하다.

 

이를 해결하기 위해서는 문제를 단순화 하면 된다.

i~j의 구간합은 다음과 같이 계산된다.

auto temp = sum[j] - sum[i -1]

이대 temp가 K와 같으면 되는 것이다.

이를 다르게 생각해 보면 sum[j]에서 K를 뺀 것이 sum[i-1]이면 된다.

즉, 우리는 sum[j - K]의 개수를 알고 있다면 j를 끝으로 하는 구간의 합이 K가 되는 개수를 구할 수 있다는 것이다.

그렇다면 누적합의 종류의 개수를 세어 기록하면서 진행한다면 sum[j - K]의 개수를 쉽게 구할 수 있다는 뜻이다.

 

코드에서는 한 개의 변수만 필요하기 때문에 j를 i로 바꿔 진행해도 된다.

map<int, long long> m;
vector<int> sum;

for (int i = 1; i <= N; i++)
{
	int x;
	cin >> x;

	sum[i] = sum[i - 1] + x;

	ans += m[sum[i] - K];
	m[sum[i]]++;
}

이때 sum[i]가 만약 K와 같다면 배열의 시작부터 i번째까지의 합이 K인 것이기 때문에 이는 유효한 경우가 된다.

이를 위해 m[0] = 1로 설정하여 정답에 더해질 수 있도록 하자.

m[0] = 1;

---

 

 

 

 

 

전체 코드

#include <bits/stdc++.h>
#include <unordered_set>
using namespace std;
#define INPUT_OPTIMIZE cin.tie(NULL); cout.tie(NULL); ios::sync_with_stdio(false);
#define INF 2e9

using namespace std;
int N, K;
map<int, long long> m;
vector<int> sum;

int main()
{
	INPUT_OPTIMIZE;
	
	cin >> N >> K;
	sum.resize(N + 1, 0);

	long long ans = 0;
	m[0] = 1;
	for (int i = 1; i <= N; i++)
	{
		int x;
		cin >> x;

		sum[i] = sum[i - 1] + x;

		ans += m[sum[i] - K];
		m[sum[i]]++;
	}

	cout << ans;

	return 0;
}

---