P6834 [Cnoi2020] BIT

题意

传送门 P6834 [Cnoi2020]梦原

题解

考虑已知树的形状，所需魔法的最少次数；即使用最少的各节点权值为 $1$ 的连通块，使各连通块点权求和后，满足树上任一节点 $i$ 点权为 $A_i$。

先考虑在链上的情况，答案为 $\sum \max (0,A_{i+1}-A_i)$。容易证明，若后继节点 $i+1$ 与当前节点 $i$ 差分为正，则 $i+1$ 需要叠加 $A_{i+1}-A_i$ 的连通分量；反之，则无需新增连通分量。同理，若在树上，答案为 $\sum \max (0,A_u-A_{parent(u)})$。那么可以依次计算各节点对答案的贡献。

设 $i+1$ 可选择的父节点数量为 $num=i-\max (-1,i-k)$，那么 $i+1$ 的贡献为 $$\frac{1}{num}\sum _{j\in (i-k,i]\cap N^{+}}\max (0,A_{i+1}-A_j)$$ 那么使用一个 $BIT$ 维护基于值域的数量和，一个 $BIT$ 维护基于值域的值域和。由于 $A$ 值域较大，进行离散化处理。总时间复杂度 $O(N\log N)$。

#include <bits/stdc++.h>
using namespace std;
#define rep(i, l, r) for (int i = l, _ = r; i < _; ++i)
typedef long long ll;
const int maxn = 1000005, mod = 998244353;
int N, K, A[maxn];
int nsx, sx[maxn], Num[maxn], Sum[maxn];

int rd()
{
    
    
    int x = 0;
    char c = 0;
    for (; !isdigit(c); c = getchar())
        ;
    for (; isdigit(c); c = getchar())
        x = (x << 1) + (x << 3) + c - '0';
    return x;
}

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

void add(int bit[], int i, int x)
{
    
    
    while (i <= N)
        bit[i] = (bit[i] + x) % mod, i += i & -i;
}

int sum(int bit[], int i)
{
    
    
    int s = 0;
    while (i)
        s = (s + bit[i]) % mod, i -= i & -i;
    return s;
}

int cps(int x) {
    
     return lower_bound(sx, sx + nsx, x) - sx; }

int main()
{
    
    
    N = rd(), K = rd();
    rep(i, 0, N) A[i] = rd(), sx[i] = A[i];
    sort(sx, sx + N);
    nsx = unique(sx, sx + N) - sx;
    rep(i, 0, N) A[i] = cps(A[i]);
    int res = sx[A[0]];
    rep(i, 0, N - 1)
    {
    
    
        add(Num, A[i] + 1, 1), add(Sum, A[i] + 1, sx[A[i]]);
        if (i - K >= 0)
            add(Num, A[i - K] + 1, -1), add(Sum, A[i - K] + 1, -sx[A[i - K]]);
        int num = i - max(-1, i - K);
        int d = ((ll)sx[A[i + 1]] * sum(Num, A[i + 1] + 1) - sum(Sum, A[i + 1] + 1)) % mod;
        d = (ll)d * pow_mod(num, mod - 2) % mod;
        res = (res + d) % mod;
    }
    printf("%d\n", (res + mod) % mod);
    return 0;
}