DOJ Array and Counting 2 c++

https://dojoi.xyz/ko/problems/85

#85 Array and Counting 2

 

문제

 

양의 정수 N과 길이가 N인 수열 A=[A1​,A2​,⋯,AN​]이 주어진다. 다음 조건을 만족하는 정수 순서쌍 (l,r)(l,r)의 개수를 구하여라.

• 1≤ l ≤ r ≤ N
• 부분 수열 [Al​,Al+1​,⋯,Ar​]의 산술평균이 정수이다.
• N은 200000이하의 정수, 1<= Ai <= 100 

 

 

접근 & 풀이

처음에는 산술평균이 정수라는 조건이 크게 와닿지 않았다.

그리고 Ai의 제한은 왜 그리 작은지도 이해가 가지 않았다. 중앙값으로 뭔가를 해야 하나?

 

아니면 누적합과 구간 길이를 이용해서 뭔가 수학적으로 풀어야 하나?? 와 같은 생각을 했다.

 

그리고 일주일이 지나 평균을 미리 정해놓은 후, 원래 수열에서 평균 만큼 뺐을 때 합이 0이 되는 구간의 개수를 누적합으로 구해주면 된다는 신기한 풀이를 들었다. 언제인가 '누적합을 이용한 합이 0인 구간 찾기'라는 개념을 배웠던 적이 있던 것 같긴 한데...

 

기억이 가물가물하니 이번에 다시 정리해보겠다.

 

 

 

i < j 일 때 

sum[j] = sum[i]를 만족시킨다고 해보자.

이때,  Ai+1 + Ai+2 +Ai+3......Aj 의 값은 0이다.

 

그런데, i 보다 작은 인덱스에서 값이 sum[j]인 것이 몇개 더 존재할 수 있다. 

때문에, 지금까지 sum의 값이 같은 것들의 개수를 저장해놓은 후, 다시 등장할 때 마다 +1을 해주면, 

누적합의 처음값과 끝 값이 같은 구간(쌍)의 개수를 O(n)에 빠르게 구할 수 있는 것이다.

 

sum의 값은 매우 커질 수 있으므로 map으로 저장하면 된다. (웬만하면 속도가 빠른 unordered map을 사용하자.)

 

array

0

-1

0

1

0

1

-1

 

sum

0

-1

-1

0

0

1

0

위와 같은 배열을 예로 들어보겠다.

합이 0인 것의 개수는 1로 초기화해두겠다.

 

map[0] = 1

cnt (합이 0인 구간의 개수) = 0

 

첫번째 누적합은 합이 0이므로 cnt += 1

map[0]++

 

두번째 누적합은 합이 -1이다. 지금까지 등장한 적이 없다. continue

map[-1]++

 

세번째 누적합도 -1이다. -1은 한번 등장했으므로 cnt += 1

map[-1]++

 

 

네번째 누적합은 다시 0이다. 0은 지금까지 2번 나왔다. cnt += 2

map[0]++ (2-> 3)

 

(후략) 

 

이와 같이 처음과 끝이 같은 값으로 짝지어질 수 있는 것의 개수를 메모이제이션을 통해 구할 수 있다.

 

 

정답코드

 

#include <bits/stdc++.h>
#define int long long
using namespace std;
int arr[202020];
int32_t main() {
	int N, ans = 0;
	cin >> N;
	for(int i = 1; i <= N; i++) {
		cin >> arr[i];
	}
	for(int i = 1; i <= 100; i++) {
		unordered_map<int, int>mp;
		int sum = 0;
		mp[0] = 1;
		for(int j = 1; j <= N; j++) {
			sum += arr[j] - i;
			if(mp[sum]) {
				ans += mp[sum];
			}
			mp[sum]++;
		}
	}
	cout << ans;
}

 

 

평균을 고정해놓는다는 것이 참신한 문제였다. 이런식으로 어떤 문제를 해결하기 쉬운 문제로 바꾸는 것이 중요한 것 같다.