Codeforces Round #400 (Div. 1 + Div. 2, combined) C. Molly's Chemicals

http://codeforces.com/contest/776/problem/C

1 ≤ n ≤ 105, 1 ≤ |k| ≤ 10

sequence 의 길이 n과 k가 입력으로 주어지고 n개의 원소가 입력으로 주어질때

원소의 합이 k^x 의 형태인 non-empty, continuous subsequence 의 개수를 찾는 문제

우선 모든 경우를 다 따져준다면 O(n^2) 로 TLE 가 난다.

우선 누적합 pSum 을 모든 인덱스에 대해 다 구해준다. 그러면

[l,r] 의 subsequence 의 원소의 합은 pSum[r]-pSum[l-1] 로 나타낼수있다. 즉

pSum[r] - pSum[l-1] = k^x 인 (l,r) 쌍의 수를 찾으면 되는데 

만약 위의 식이 성립한다면 항을 적절히 넘겨서 이런식을 도출해 낼 수가 있다.

pSum[r] - k^x = pSum[l-1] 

즉 위의 식을 만족하는 (l,r) 쌍을 찾으면 되는데, r 이 고정일 때 가능한 k의 거듭제곱들을 모두 빼 봐서 r보다 작은 인덱스를 가지는 pSum이 존재한다면 그 수 만큼 답을 누적시켜주면 되는것이다.

이 것을 하기 위해

우선 첫번째로, 가능한 k의 거듭제곱들을 모두 구하고

두번째로, pSum 배열을 인덱스 i를 1부터 차례로 보면서 r을 이 인덱스로 고정시키고 위에서 말했던 것처럼 가능한 k의 거듭제곱들을 pSum[i]에서 빼봐서 앞에서 나온 인덱스에대한 pSum 값이 존재했었는지를 확인하여 답을 갱신시켜준다.

우선 첫번째걸 해보자 가능한 k의 거듭제곱은 모든 원소의 합의 절대값이 10^14 보다 작거나같고 아무리 많아봐야 k가 2일때이므로 한 50개도 안될것이다. 그냥 구하면되고

두번째걸 해보자. 앞에 나왔던 pSum의 개수를 세어줘야하므로 각각의 반복문에서 pSum값을 인덱스로 하고 개수를 값으로 갖는 배열을 만들어보려고하니 pSum이 최대10^14 이기때문에 공간이 부족하다. 값을 인덱스로 하고싶을때 총 원소의 수는 적은데 인덱스가 너무 커서 배열을 사용하지 못할 때 유용한 자료구조가 있다. map이다. map을 사용하면 key값을 인덱스로 삼아 접근을 할 수가있기 때문에 마치 배열을 잡은 것 처럼 쓸수가있다.

매 반복마다 k의 가능한 거듭제곱을 pSum[i]에서 뺀값을 인덱스로 하는 값을 답에 누적시키고 map[pSum[i]]++ 를 해주면된다.

#include <bits/stdc++.h>
#define REP(i,a,b) for(int i=a;i<=b;++i)
#define FOR(i,n) for(int i=0;i<n;++i)
#define pb push_back
#define all(v) (v).begin(),(v).end()
#define sz(v) ((int)(v).size())
#define inp1(a) scanf("%d",&a)
#define inp2(a,b) scanf("%d%d",&a,&b)
#define inp3(a,b,c) scanf("%d%d%d",&a,&b,&c)
#define inp4(a,b,c,d) scanf("%d%d%d%d",&a,&b,&c,&d)
#define inp5(a,b,c,d,e) scanf("%d%d%d%d%d",&a,&b,&c,&d,&e)
using namespace std;
typedef long long ll;
typedef pair<ll,ll> pll;
typedef vector<int> vi;    
typedef vector<ll> vl;
typedef pair<int,int> pii;
typedef vector<pii> vii;
typedef vector<pll> vll;
typedef vector<vector<int> > vvi;
typedef pair<int,pair<int,int> > piii;
typedef vector<piii> viii;
const double EPSILON = 1e-9;
const double PI = acos(-1);
const int MOD = 1e9+7;
const int INF = 0x3c3c3c3c;
const long long INFL = 0x3c3c3c3c3c3c3c3c;
const int MAX_N = 102;
map<ll, ll> um;
int N,K;
int A[100003];
ll pSum[100003];
unordered_set<ll> us;
vector<ll> cand;
int main() {
    inp2(N,K);
    FOR(i,N) inp1(A[i]), pSum[i+1]=pSum[i]+A[i];
    long long pusher=1;
    for(; abs(pusher) <= 1e16; pusher*=K) {
        if(us.count(pusher)) break;
        us.insert(pusher);
        cand.pb(pusher);
    }
    ll ans=0;
    um[0]=1;
    REP(i,1,N){
        for(auto kk : cand) ans+=um[pSum[i]-kk];
        um[pSum[i]]++;
    }
    printf("%lld",ans);
    return 0;
}