Codeforces Round 955 E. Number of k-good subarrays【分治、记忆化】

E. Number of k-good subarrays

题意

定义 $bi t (x)$ 为 $x$ 的二进制表示下 $1$ 的数量
一个数组的子段被称为 $k - g oo d$ 的当且仅当：对于这个子段内的每个数 $x$ ，都有 $\leq k$

现在给定一个长度为 $n$ 的数组 $a$ ，且 $\forall 0 \leq i < n, a_i = i$ ，即 $a = \{0, 1, 2, ... n - 1\}$

求出 $a$ 中 $k - g oo d$ 的子段数量

思路

首先我们观察发现：可以先将原数组选出连续的若干段，使得每一段都是极大的 $k - g oo d$ 子段，也就是不能再往两边扩了。假设这一段的长度为 $l e n$ ，那么这一段对答案的贡献是： $C_{len+1} ^ 2$

由于问题规模过大，我们考虑分治：
定义 $f (n, k)$ 为在 $[0, n - 1]$ 中的 $k - g oo d$ 子段数量，这个子状态返回一个元组 ${res, l, r\}$

$res$ 表示 $[0, n - 1]$ 里符合条件的子段数量
$l$ 表示从头开始的最长连续 $k - g oo d$ 长度
$r$ 表示从尾开始的最长连续 $k - g oo d$ 长度

我们再找到一个最大的 $m$ ，使得 $2^m < n$ ，即 $\leq n - 1$
并将 $[0, n - 1]$ 分成两段： $0, 2 ^ m - 1]$ 和 $2^m, n - 1]$
我们发现第一段其实就是子问题 $f(2 ^ m, k)$ ，那么第二段呢？它并不从零开始，好像无法转换成某一个子状态。

__builtin_clzll ( n ) 可以算出 $n$ 的前导 $0$ 数量

其实 $2^m$ 的最高位一定是和 $n - 1$ 的最高位相同的，那么我们就等价于提前为第 $m$ 为预定了一个 $1$ ，第二段就转换成了子状态： $f(n - 2^m, k - 1)$

因此， $f(n, k) = f(2 ^ m, k) + f(n - 2 ^ m, k - 1)$

不难发现，第二部分的复杂度为 $O(\log n)$ ，而第一部分由于它的第一个分量永远为 $2$ 的整数次幂，因此第一部分最多有 $\cdot \log n)$ 种状态，我们可以使用记忆化加速搜索

当我们得到两个部分的返回 ${res_1, l_1, r_1 \}$ 和 ${res_2, l_2, r_2 \}$ 之后，我们可以归并答案，由于之前我们把 $r_1$ 向前延伸和 $l_2$ 向后延伸的算在了 $res_1$ 和 $res_2$ 里，因此我们在合并中间的连续段的时候（合二为一），需要减掉之前的错误统计。

即 $res = res_1 + res_2 - C_{r_1 + 1}^2 - C_{l_2 + 1}^2 + C_{r_1 + l_2 + 1}^2$

#include<bits/stdc++.h>
#define fore(i,l,r)	for(int i=(int)(l);i<(int)(r);++i)
#define fi first
#define se second
#define endl '\n'
#define ull unsigned long long
#define ALL(v) v.begin(), v.end()
#define Debug(x, ed) std::cerr << #x << " = " << x << ed;

const int INF=0x3f3f3f3f;
const long long INFLL=1e18;

typedef long long ll;

template<class T>
constexpr T power(T a, ll b){
    T res = 1;
    while(b){
        if(b&1) res = res * a;
        a = a * a;
        b >>= 1;
    }
    return res;
}

constexpr ll mul(ll a,ll b,ll mod){ //快速乘，避免两个long long相乘取模溢出
    ll res = a * b - ll(1.L * a * b / mod) * mod;
    res %= mod;
    if(res < 0) res += mod; //误差
    return res;
}

template<ll P>
struct MLL{
    ll x;
    constexpr MLL() = default;
    constexpr MLL(ll x) : x(norm(x % getMod())) {}

    static ll Mod;
    constexpr static ll getMod(){
       if(P > 0) return P;
       return Mod;
    }

    constexpr static void setMod(int _Mod){
       Mod = _Mod;
    }
    constexpr ll norm(ll x) const{
       if(x < 0){
           x += getMod();
       }
       if(x >= getMod()){
           x -= getMod();
       }
       return x;
    }
    constexpr ll val() const{
       return x;
    }
    explicit constexpr operator ll() const{ 
       return x; //将结构体显示转换为ll类型： ll res = static_cast<ll>(OBJ)
    }
    constexpr MLL operator -() const{ //负号，等价于加上Mod
       MLL res;
       res.x = norm(getMod() - x);
       return res;
    }
    constexpr MLL inv() const{
       assert(x != 0);
       return power(*this, getMod() - 2); //用费马小定理求逆
    }
    constexpr MLL& operator *= (MLL rhs) & { //& 表示“this”指针不能指向一个临时对象或const对象
       x = mul(x, rhs.x, getMod()); //该函数只能被一个左值调用
       return *this;
    }
    constexpr MLL& operator += (MLL rhs) & {
       x = norm(x + rhs.x);
       return *this;
    }
    constexpr MLL& operator -= (MLL rhs) & {
       x = norm(x - rhs.x);
       return *this;
    }
    constexpr MLL& operator /= (MLL rhs) & {
       return *this *= rhs.inv();
    }
    friend constexpr MLL operator * (MLL lhs, MLL rhs){
       MLL res = lhs;
       res *= rhs;
       return res;
    }
    friend constexpr MLL operator + (MLL lhs, MLL rhs){
       MLL res = lhs;
       res += rhs;
       return res;
    }
    friend constexpr MLL operator - (MLL lhs, MLL rhs){
       MLL res = lhs;
       res -= rhs;
       return res;
    }
    friend constexpr MLL operator / (MLL lhs, MLL rhs){
       MLL res = lhs;
       res /= rhs;
       return res;
    }
    friend constexpr std::istream& operator >> (std::istream& is, MLL& a){
       ll v;
       is >> v;
       a = MLL(v);
       return is;
    }
    friend constexpr std::ostream& operator << (std::ostream& os, MLL& a){
       return os << a.val();
    }
    friend constexpr bool operator == (MLL lhs, MLL rhs){
       return lhs.val() == rhs.val();
    }
    friend constexpr bool operator != (MLL lhs, MLL rhs){
       return lhs.val() != rhs.val();
    }
};

const ll mod = 1e9 + 7;
using Z = MLL<mod>;

const ll INV = 500000004ll; //2的逆元

std::map<std::pair<ll, int>, std::tuple<Z, ll, ll>> mem; //记忆化

std::tuple<Z, ll, ll> dfs(ll n, int k){
    if(k < 0)   return {0, 0, 0};
    if(n == 1)  return {1, 1, 1}; //小于1只有0
    int m = 63 - __builtin_clzll(n);
    ll mid = 1ll << m;
    if(mid == n){
        --m;
        mid >>= 1;
        if(mem.count({n, k})) return mem[{n, k}];
    }

    auto [res1, l1, r1] = dfs(mid, k);
    auto [res2, l2, r2] = dfs(n - mid, k - 1);

    Z w1 = Z(r1) * (r1 + 1) * INV;
    Z w2 = Z(l2) * (l2 + 1) * INV;
    Z tmp(r1 + l2);
    Z res = res1 + res2 - w1 - w2 + tmp * (tmp + 1) * INV;
    ll l = l1, r = r2;
    if(l1 == mid) l += l2;
    if(r2 == n - mid)   r += r1;

    if((mid << 1) == n) mem[{n, k}] = {res, l, r};

    return {res, l, r};
}

int main(){
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    std::cout.tie(nullptr);
    int t;
    std::cin >> t;
    while(t--){
        ll n;
        int k;
        std::cin >> n >> k;
        auto [ans, l, r] = dfs(n, k);
        std::cout << ans << endl;
    }
    
    return 0;
}