Luogu P6834 梦原 题解

原题传送门

首先考虑如果树的形态确定了之后的情况：如果当前点比爸爸的值大，那么显然爸爸变成 \(0\) 之后，这个点需要自己被额外操作删除，贡献就是 \(a[v] - a[u]\)。
类似，如果比爸爸的值小，那么这个点肯定会跟爸爸一起被删除，所以贡献就是 \(0\)。
综上所述，一个点的贡献是 \(max{a[v] - a[u], 0}\)。

接下来考虑树的形态，以求得答案。
首先，最终期望值可以拆分成每个点的贡献的期望值之和，那么显然有动态规划：

\(dp[i] = \frac{1}{i - max(1, i - k) } * \sum_{v \in [max(1, i - k), i - 1]}{max(a[v] - a[u], 0)}\)

其中 \(dp[1] = a[1]\)

答案就是 \(\sum_{i \in [1, n]}{dp[i]}\)

注意优化：使用权值树状数组快速求得\(a_u\)总和，然后再用一权值树状数组快速求具体数量，乘上 \(a_v\) 即可。注意划窗时，要删去左侧的部分。

#include <bits/stdc++.h>
using namespace std;
#define N 2000010
#define ll long long
#define int long long

template <class T>
inline T read(T& a){
    T x = 0, s = 1;
    char c = getchar(); 
    while(!isdigit(c)){ if(c == '-')  s = -1; c = getchar(); }
    while(isdigit(c)){ x = x * 10 + (c ^ '0'); c = getchar(); }
    a = x * s;
    return a; 
}

int n, k; 
ll a[N]; 
const int mod = 998244353; 
ll dp[N]; 

ll qpow(ll a, ll b){
    ll sum = 1; 
    while(b){
        if(b & 1) sum = (sum * a) % mod; 
        b >>= 1; 
        a = (a * a) % mod; 
    }
    return sum; 
}

unordered_map <int, int> g; 
int b[N]; 

struct Segment_tree{
    int tree[N]; 

    #define lowbit(x) (x&-x)

    void update(int x, int k){
        for(int i = x; i <= n; i += lowbit(i))
            tree[i] = ((tree[i] + k) % mod + mod) % mod;
        return ;  
    }

    int ask(int x){
        int sum = 0; 
        for(int i = x; i; i -= lowbit(i)){
            sum = (sum + tree[i]) % mod; 
        }
        return sum; 
    }

    int query(int l, int r){
        return (ask(r) - ask(l - 1) + mod) % mod; 
    }
 
} tree, tree2; 

signed main(){
    read(n), read(k); 
    for(int i = 1; i <= n; i++)
        b[i] = read(a[i]);
    int id = 0; 
    sort(b + 1, b + n + 1); 
    for(int i = 1; i <= n; i++){
        if(!g[b[i]]) g[b[i]] = ++id; 
    }
    for(int i = 1; i <= n; i++){
        b[i] = g[a[i]]; 
    }
    // cout << qpow(2, mod - 2) << endl; 
    dp[1] = a[1]; 
    tree.update(b[1], a[1]); 
    tree2.update(b[1], 1); 
    int now = 1; 
    for(int i = 2; i <= n; i++){
        while(now < max(1ll, i - k)) tree.update(b[now], -a[now]), tree2.update(b[now], -1), now++; 
        // cout << "------" << -tree.query(1, 1, n, 1, g[a[i]] - 1) + a[i] * tree2.query(1, 1, n, 1, g[a[i]] - 1) << " " << qpow(i - max(1, i - k), mod - 2) << endl; 
        dp[i] = (qpow(i - max(1ll, i - k), mod - 2) % mod * ((-tree.query(1, b[i] - 1) + (a[i] * tree2.query(1, b[i] - 1)) % mod) % mod)) % mod; 
        tree.update(b[i], a[i] % mod); 
        tree2.update(b[i], 1); 
    }
    // for(int i = 1; i <= n; i++)
        // cout << dp[i] << endl; 
    ll ans = 0; 
    for(int i = 1; i <= n; i++){
        ans = (ans + dp[i]) % mod; 
    }

    cout << (ans + mod) % mod << endl; 
    return 0; 
}