题解：CF1982E Number of k-good subarrays

SFlyer · 2024-06-27 15:46:31 · 题解

设 sol(n,k) 为答案。（注意，是 0\sim n-1 的！）那么，我们 sol(n,k) 可以通过 sol(mx,k) 和 sol(n-mx,k-1) 算出来，其中 mx 是小于 n 的最大 2 次幂。（很好理解，大于等于 mx 的就会有一位固然是 1，就会是 k-1。）设 mx=2^c。

问题是怎么算。除了内部的，还有一个端点在左边一个端点在右边的。官方题解是维护了三个值，很复杂。有没有直接算的方法？

我们发现，如果 c>k，那么左边最后一个就会有大于 k 个 1，没有贡献，所以必须 c\le k。这个时候左边所有的都满足，只需要计算右边的。右边的到了 2^k-1 也不行，因为这样也超过了（当然要和长度取最小值）。所以我们算出了右边的贡献 s=\min(2^c-1,n-2^c)。那么，答案就要多加上 s\cdot 2^c。

直接记忆化搜索即可，复杂度是 \mathcal{O}(k \log n) 的。代码非常短。

#include <bits/stdc++.h>

using namespace std;

using ll = long long;

const ll mod = 1e9+7;

map<pair<ll,ll>,ll> mp;

ll sol(ll n,ll k){
    if (k==0 || n==1){
        return 1;
    }
    if (mp.count({n,k})){
        return mp[{n,k}];
    }
    ll c=0;
    while ((1ll<<c+1)<n){
        c++;
    }
    ll l=sol(1ll<<c,k),r=sol(n-(1ll<<c),k-1);
    ll ans=l+r;
    if (c<=k){
        ll s=min((1ll<<k)-1,n-(1ll<<c));
        ans+=s%mod*((1ll<<c)%mod)%mod;
    }
    mp[{n,k}]=ans%mod;
    return ans%mod;
}

void solve(){
    ll n,k;
    cin>>n>>k;
    cout<<sol(n,k)<<"\n";
}

int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);

    int t;
    cin>>t;
    while (t--){
        solve();
    }
    return 0;
}