P6834 [Cnoi2020]梦原 树状数组 期望DP

P6834 [Cnoi2020]梦原

题目链接

​ 树状数组优化期望DP。

​ 我们可以发现，当\(k= 1\)时，这道题就和积木大赛一样，有一个非常有用的结论：对于每个点，只要把\(max(0, a[i] - a[i - 1])\)累加起来就是最终答案。假设前一个比后一个高，那么前一个为0时后一个一定早就为0了；假设后一个比前一个高，那么后一个多使用的操作就是后一个的高度减去前一个的高度。

​ 这个题把问题转化到了树上。因为\(k\)是不确定的，又因为要求期望，所以我们可以写出一个DP转移方程：\(f[i] = \displaystyle P(min(k, i - 1))\sum _{j = i - k}^{i - 1} a[i] - a[j]\)。这个式子应该挺好理解，就是\(E = \displaystyle \sum P *V\)。因为每个\(P\)都一样，所以可以提出来。

​ 暴力算肯定是不行的，所以我们用树状数组维护一下。维护一个前缀和，再维护一个个数的前缀和，上式就可以化成\(cnt * a[i] - \sum a[j]\)。

#include <bits/stdc++.h>

#define int long long

using namespace std;

inline long long read() {
    long long s = 0, f = 1; char ch;
    while(!isdigit(ch = getchar())) (ch == '-') && (f = -f);
    for(s = ch ^ 48;isdigit(ch = getchar()); s = (s << 1) + (s << 3) + (ch ^ 48));
    return s * f;
}

const int N = 1e6 + 5, mod = 998244353;
int n, k, ans;
int a[N], b[N], c[N], t_num[N], t_sum[N];

int lowbit(int x) { return x & -x; } 

void change(int x, int sum, int num) { for(; x < N; x += lowbit(x)) t_sum[x] += sum, t_num[x] += num; }

int query_sum(int x) { int res = 0; for(; x ; x -= lowbit(x)) (res += t_sum[x]) %= mod; return res; }

int query_num(int x) { int res = 0; for(; x ; x -= lowbit(x)) (res += t_num[x]) %= mod; return res; }

int ksm(int x, int y) {
    int res = 1;
    while(y) {
        if(y & 1) res = 1ll * res * x % mod;
        x = 1ll * x * x % mod; y >>= 1; 
    }
    return res;
}

int inv(int x) {
    return ksm(x, mod - 2);
}

signed main() {

    n = read(); k = read();
    for(int i = 1;i <= n; i++) b[i] = a[i] = c[i] = read();
    
    sort(b + 1, b + n + 1);
    int cnt = unique(b + 1, b + n + 1) - b - 1;
    for(int i = 1;i <= n; i++) a[i] = lower_bound(b + 1, b + cnt + 1, a[i]) - b; //离散化

    ans = c[1]; change(a[1], c[1], 1);
    for(int i = 2;i <= n; i++) {
        if(i >= k + 2) change(a[i - k - 1], -c[i - k - 1], -1); //不在范围内的减去
        (ans += 1ll * inv(i <= k + 1 ? i - 1 : k) * ( 1ll * query_num(a[i]) * c[i] % mod - query_sum(a[i])) % mod) %= mod; 
        (ans += mod) %= mod;
        change(a[i], c[i], 1);
    }

    printf("%lld", ans);

    return 0;
}