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

posted on 2024-02-22 06:04:01 | under 未分类 | source

考虑合法区间的判定。先单调栈求出 \(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;
}