백준 25332 / C++

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

25332번: 수들의 합 8

 

[ 풀이 ] 

 

B[i]-A[i]=C[i]라고 하자.  C[i]+...+C[j]=0 인 (i,j) 쌍을 구하자. 

C[i]까지 누적합을 pref[i]라 하면, pref[i-1]==pref[j]인 (i,j)를 구하면 된다. 

단순하게 O(N^2)에 구하면 TLE이다. 

값과 횟수를저장하는 배열 X를 만들어주면 쉽다. (X[값]=횟수) 

만약 지금 pref[i]의 값이 3이 나왔고, 다음에 X[pref[j]]가 0이 아니라면(즉, 이미 갱신된적 있다.)

값 X[pref[j]]는 지금까지 pref[i]=3이 나온 횟수가 된다. 따라서 j미만 모든 쌍을 셀 수 있다.

(i,j)를 센 이후엔 X[pref[j]]값을 1 더해주면 갱신된다. 

이러면 O(N)에 해결할 수 있다. 

그런데 문제는 pref[i]가 너무 커질 수 있다. C[i]는 최대 10^4이고,  N<=2e5이므로, 

pref는 2e9까지 나와서 배열에 담는건 불가능하다. 또한 pref값은 산발적이므로, 

모든 값 X[0]~X[2e9]를 알 필요는 없다. 이럴때 map을 써주면 된다. 

map은 그 value와 key만을 저장하므로 pref값이 커져도 가능하다.

 

[ 코드 ] 

 

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
ll A[200001], B[200001], pref[200001];
map<ll, ll> X;
#define for(i,n) for (int i=1;i<=n;i++)
 
int main() {
    ios::sync_with_stdio(0); cin.tie(0);
    int n; cin >> n;
    for (i, n) cin >> A[i];
    for (i, n) cin >> B[i];
    // 누적합 만들기
    for (i, n) pref[i] = pref[i - 1] + B[i] - A[i];
    ll ans = 0;
    // pref[i-1]=pref[j]인 (i,j) 세기
    for (i, n) {
        if (pref[i] == 0) ans++; // i=1인 경우 pref[i]=0인걸 모두 세준다.(예외처리) 
        ans += X[pref[i]]++; // (... , i) 쌍 모두 세주고 갱신 
    }
    cout << ans << '\n';
}