Codeforces Round #400 C. Molly's Chemicals （二分）

题意

给定 n 个数的数组 $a_1, a_2, ..., a_n (10^{-9}\le a_i\le 10^9)$ 。要求找出有多少个区间 $[l, r]$ ，使得区间内的所有数的和恰好是 k 的非负幂次。其中 $1\le n\le10^5, \ 1\le\mid k \mid \le 10$

分析

预处理 $sum[i] = \Sigma_{j=1}^i{a_j}$ 。对于每个位置 i， 找到只需要找的每个值为 $sum[i] + k^{pwr}, 0\le pwr \le \log_2{(10^9\times 10^5)}$ 的位置 j 。通过一个二元组 $pair<long\ long, int> p[N]$ 暂存，排序保证$p[i].first < p[i+1].first$，每次二分查找即可。

代码

#include<bits/stdc++.h>
using namespace std;
const long long inf = 1e14;
const int N = 100000 + 10;
long long a[N], sum[N];
pair<long long, int> p[N];
bool operator<(pair<long long, int> a, pair<long long, int> b) {
    if(a.first == b.first)
        return a.second < b.second;
    return a.first < b.first;
}
int main()
{
    int n, k;
    scanf("%d %d",&n,&k);
    for(int i=1;i<=n;i++)
    {
        scanf("%I64d",&a[i]),
        sum[i] = sum[i-1] + a[i],
        p[i].first = sum[i],
        p[i].second = i;
    }
    sort(p+1, p+n+1);
    long long ans = 0;
    for(int i=0;i<=n;i++)
    {
        long long val = 1,  aim;
        for(int pwr=0;val <= inf;val*=k, pwr++)
        {
            if(k == 1 && pwr == 1)  break;
            if(k == -1 && pwr == 2) break;
            aim = sum[i] + val;
            int l = lower_bound(p+1, p+n+1, make_pair(aim, i)) - p;
            int r = upper_bound(p+1, p+n+1, make_pair(aim+1, 0)) - p;
            ans += (r-l);
        }
    }
    printf("%I64d\n",ans);
}