atc abc464F 思路分享（期望，meet-in-the-middle）

https://atcoder.jp/contests/abc464/tasks/abc464_f

题意

有 \(n\) 个保险柜，第 \(i\) 个保险柜有 \(a_i\) 元，小偷每次随机打开一个保险柜取走里面所有的钱，直到取走钱总和 \(\ge X\)，求小偷取走钱总和的期望，模 \(998244353\).

\(1\le n \le 40\).

思路

根据期望的线性性，\(\mathbb{E} = \sum_{i=1}^{n}{a_iP_i}\)，其中 \(P_i\) 是保险柜 \(i\) 被打开的概率.

考虑求 \(P_i\)，枚举 \(S\) 表示在 \(i\) 之前被打开的保险柜集合，在 \(i\) 之前，\(S\) 中元素可以任意排列；在 \(i\) 之后，剩余元素可以任意排列，即：

\[P_i = \sum_{S}{\frac{|S|! \cdot (n-1-|S|)!}{n!}} \]

其中 \(S\) 必须满足 \(\sum_{i\in S}{a_i} \lt X\)，且 \(i\notin S\).

因此：

\[\mathbb{E} = \sum_{i=1}^{n}{a_i \sum_{S}{\frac{|S|!\cdot (n-1-|S|)!}{n!}}} \]

交换求和：

\[\mathbb{E} = \sum_{S}{\frac{|S|!\cdot (n-1-|S|)!}{n!}\sum_{i\notin S}{a_i}} \]

令 \(total = \sum_{i=1}^{n}{a_i}\)，则 \(\sum_{i\notin S}{a_i} = total-\sum_{i\in S}{a_i}\).

直接枚举 \(S\) 是困难的，考虑固定 \(k=|S|\) 批量求和，令 \(C_k\) 表示 \(|S|=k\) 的数量，\(T_k\) 表示 \(|S|=k\) 的 \(\sum_{i\in S}{a_i}\) 之和，原式转化成：

\[\mathbb{E} = \sum_{k}{\frac{k! \cdot (n-1-k)!}{n!} \cdot (total\cdot C_k-T_k)} \]

问题转化成求 \(C\) 和 \(T\).

注意到 \(n\le 40\) 是可以折半的范围，考虑 meet-in-the-middle.

对于左半部分，直接搜出所有集合，维护子集大小 \(cnt\) 和 子集元素和 \(sum\).

对于右半部分，按子集大小分类，\(right_k\) 存所有子集大小为 \(k\) 的子集元素和，升序排序，并预处理前缀和.

枚举左半部分集合，在右半部分的长度桶内二分找到最后一个 \(pos\) 满足 \(sum+right_{pos} \le X\) 的位置，根据前缀和批量地更新 \(C\) 和 \(T\) 即可.

时间复杂度 \(\mathcal{O}(n\log n \cdot 2^{n/2})\).

代码

//author:kzssCCC

#include <bits/stdc++.h>
using namespace std;
using ll = long long;

template<int MOD>
struct modint {
    int val;
    
    modint() : val(0) {}
    modint(long long v) {
        val = v % MOD;
        if (val < 0) val += MOD;
    }
    
    modint& operator++() { val = (val + 1 == MOD ? 0 : val + 1); return *this; }
    modint& operator--() { val = (val == 0 ? MOD - 1 : val - 1); return *this; }
    modint operator++(int) { modint res = *this; ++*this; return res; }
    modint operator--(int) { modint res = *this; --*this; return res; }
    
    modint& operator+=(const modint& o) { val += o.val; if (val >= MOD) val -= MOD; return *this; }
    modint& operator-=(const modint& o) { val -= o.val; if (val < 0) val += MOD; return *this; }
    modint& operator*=(const modint& o) { val = 1LL * val * o.val % MOD; return *this; }
    modint& operator/=(const modint& o) { return *this *= o.inv(); }
    
    friend modint operator+(modint a, const modint& b) { return a += b; }
    friend modint operator-(modint a, const modint& b) { return a -= b; }
    friend modint operator*(modint a, const modint& b) { return a *= b; }
    friend modint operator/(modint a, const modint& b) { return a /= b; }

    friend bool operator==(const modint& a, const modint& b) { return a.val == b.val; }
    friend bool operator!=(const modint& a, const modint& b) { return a.val != b.val; }
    modint operator-() const { modint res = *this; res.val = (res.val == 0 ? 0 : MOD - res.val); return res;};
    modint operator+() const { return *this; };
    
    modint qpow(long long p) const {
        modint res = 1, a = *this;
        while (p > 0) {
            if (p & 1) res *= a;
            a *= a;
            p >>= 1;
        }
        return res;
    }

    modint inv() const {
        return qpow(MOD - 2);
    }
    
    friend std::ostream& operator<<(std::ostream& os, const modint& m) { return os << m.val; }
    friend std::istream& operator>>(std::istream& is, modint& m) { long long v; is >> v; m = modint(v); return is; }
};
using Z = modint<998244353>;

const int MAXN = 40;
Z fac[MAXN+1],inv[MAXN+1];

void init(){
    fac[0] = 1;
    for (int i=1;i<=MAXN;i++){
        fac[i] = fac[i-1]*i;
    }

    inv[MAXN] = fac[MAXN].inv();
    for (int i=MAXN-1;i>=0;i--){
        inv[i] = inv[i+1]*(i+1);
    }
}

void solve(){
    int n;
    ll X;
    cin >> n >> X;

    vector<ll> a(n+1);
    for (int i=1;i<=n;i++){
        cin >> a[i];
    }

    vector<pair<int,ll>> left;
    for (int u=0;u<1<<n/2;u++){
        int cnt = __builtin_popcount(u);
        ll sum = 0;
        for (int j=0;j<n/2;j++){
            if (u>>j&1){
                sum += a[j+1];  
            }
        }
        if (sum<X){
            left.emplace_back(cnt,sum);
        }
    }

    vector<vector<ll>> right((n+1)/2+1);
    for (int u=0;u<1<<(n+1)/2;u++){
        int cnt = __builtin_popcount(u);
        ll sum = 0;
        for (int j=0;j<(n+1)/2;j++){
            if (u>>j&1){
                sum += a[j+n/2+1];
            }
        }
        if (sum<X){
            right[cnt].push_back(sum);            
        }
    }

    vector<vector<Z>> pre((n+1)/2+1);
    for (int k=0;k<=(n+1)/2;k++){
        sort(right[k].begin(),right[k].end());
        int len = right[k].size();

        vector<Z> temp(len);
        for (int i=0;i<len;i++){
            temp[i] = (i-1>=0?temp[i-1]:0)+right[k][i];  
        }
        pre[k] = temp;
    }

    vector<Z> C(n+1),T(n+1);
    for (auto& [cnt,sum]:left){
        for (int k=0;k<=(n+1)/2;k++){
            int pos = lower_bound(right[k].begin(),right[k].end(),X-sum)-right[k].begin()-1;
            if (pos>=0){
                C[cnt+k] += pos+1;
                T[cnt+k] += pre[k][pos]+(Z)sum*(pos+1);
            }            
        }        
    }

    Z total = accumulate(a.begin()+1,a.end(),Z(0));
    Z res = 0;
    for (int k=0;k<n;k++){
        res += fac[k]*fac[n-1-k]*inv[n]*(total*C[k]-T[k]);
    }
    cout << res << '\n';
}

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

    int t = 1;
    // cin >> t;
    while (t--) solve();

    return 0;
}