[差分] [树上问题] P10174 Great Segments 题解

· · 题解

考虑合法区间的判定。先单调栈求出 nxt_i 表示 a_i 右边第一个大于它的元素下标，易发现在序列 b 中，从一段相同元素 [l1,r1] 跳到下一段相同元素 [l2,r2]，当且仅当 l2=nxt_{l1}。

对于一段区间 [l,r]，从 l 开始跳 nxt，即不断让 l=nxt_{l} 直到 l>r，那么依据题意该区间是合法区间需满足：1<cnt<=k，其中 cnt 是总共跳的元素个数；并且最终停在 r 上（上述跳 nxt 过程中遍历到 r）。

于是可以得到 O(n^2) 算法，暴力枚举每个合法区间并打上差分标记。

考虑优化，将差分标记拆为加法和减法，尝试对于 a_i 快速求出最终在 i 上打了多少个加法标记、在 i+1 上打上了多少个减法标记，即有多少个合法区间以 i 开头和结尾。

由于每个 i 对应唯一一个 nxt_i，考虑将 nxt_i\to i 连一条边，显然会构成森林，其中 nxt_i=0 的元素就是根。

问题被转化为树上问题，对于加标记，可以快速计算得出，记 dep_i 表示 i 的深度，根的深度为 1，则 i 上的加标记数量为 \min(dep,k)-1。

对于减标记，考虑递推求解，记 d_i 表示 i 上减标记数量，kson_i 表示 i 的所有子孙节点中满足到 i 距离（最短路径边数）为 k 的节点数量，V_i 是 i 的儿子集合。则有 d_i=(\sum \limits_{j\in V_i} d_j+1) - kson_i。

最后考虑 kson_i 的求法，只要找出每个点的 k 级祖先即可，可以用栈维护当前点的所有祖先，复杂度线性。

时间复杂度 O(n)，空间复杂度 O(n)。

代码

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

const int N = 1e6 + 5;
int n, k, a[N], stk[N], st, d[N], nxt[N];
vector<int> to[N];
long long ans[N];

inline void init(int u){
    if(st - k + 1 > 0) --d[stk[st - k + 1]];
    stk[++st] = u; int stt = st;
    for(auto v : to[u])
        st = stt, init(v);
}
inline void dfs(int u, int dep){
    ans[u] += min(dep, k) - 1;
    for(auto v : to[u]){
        dfs(v, dep + 1);
        d[u] += d[v] + 1;
    }
    ans[u + 1] -= d[u];
}
signed main(){
    cin >> n >> k;
    for(int i = 1; i <= n; ++i) scanf("%lld", &a[i]);
    for(int i = n; i >= 1; --i){
        while(st && a[stk[st]] <= a[i]) --st; 
        if(st) nxt[i] = stk[st], to[stk[st]].push_back(i);
        stk[++st] = i;
    }
    for(int i = 1; i <= n; ++i)
        if(!nxt[i]) st = 0, init(i), dfs(i, 1);
    for(int i = 1; i <= n; ++i)
        ans[i] += ans[i - 1], printf("%lld\n", ans[i]); 
    return 0;
}