题解 AVL 树

link

Description

给出一个 \(n\) 个点的 AVL 树，求保留 \(k\) 个点使得字典序最小。

\(n\le 5\times 10^5\)

Solution

因为我很 sb ，所以只会 \(\Theta(n\log^2n)\)。

首先可以注意到的是，树高是 \(\Theta(\log n)\) 的，然后我们要让字典序最小的话，可以考虑一个点一个点加进入判断是否可以。

我们考虑设 \(f_{u,i}\) 表示以 \(u\) 为根的子树在当前已选的点的情况下保留深度为 \(i\) 的还需选的最小点数。那么对于我们当前考虑的点，如果已经加入的点数加上还需加入的最小点数 \(\le k\) 那么我们就可以加入这个点。

发现这个 \(f\) 每次会改变的只会有 \(\text{rt}\to u\) 这一条路径（\(u\) 是当前考虑的点），所以我们就可以做到 \(\Theta(n\log^2 n)\) 了。

Code

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

#define Int register int
#define MAXN 500005

template <typename T> inline void read (T &t){t = 0;char c = getchar();int f = 1;while (c < '0' || c > '9'){if (c == '-') f = -f;c = getchar();}while (c >= '0' && c <= '9'){t = (t << 3) + (t << 1) + c - '0';c = getchar();} t *= f;}
template <typename T,typename ... Args> inline void read (T &t,Args&... args){read (t);read (args...);}
template <typename T> inline void write (T x){if (x < 0){x = -x;putchar ('-');}if (x > 9) write (x / 10);putchar (x % 10 + '0');}
template <typename T> inline void chkmax (T &a,T b){a = max (a,b);}
template <typename T> inline void chkmin (T &a,T b){a = min (a,b);}

int n,K,rt,num,ans[MAXN],ls[MAXN],rs[MAXN],h[MAXN],f[MAXN],hei[MAXN],dp[MAXN][25];

void dfs1 (int u){
	if (ls[u]) dfs1 (ls[u]);
	if (rs[u]) dfs1 (rs[u]);
	hei[u] = max (hei[ls[u]],hei[rs[u]]) + 1;
}

int tot,pos[25],reh[25],sta[25][25];

void ins (int root,int x){
	++ tot,pos[tot] = root,reh[tot] = h[root];
	for (Int i = 0;i <= hei[root];++ i) sta[tot][i] = dp[root][i];
	if (root == x){
		h[x] = max (h[ls[x]],h[rs[x]]) + 1,memset (dp[x],0x3f,sizeof (dp[x]));
		for (Int i = h[x];i <= hei[x];++ i)
			chkmin (dp[x][i],dp[ls[x]][i - 1] + dp[rs[x]][i - 1]),
			chkmin (dp[x][i],dp[ls[x]][i - 2] + dp[rs[x]][i - 1]),
			chkmin (dp[x][i],dp[ls[x]][i - 1] + dp[rs[x]][i - 2]);
		return ;
	}
	else{
		if (x < root) ins (ls[root],x);else ins (rs[root],x);
		h[root] = max (h[ls[root]],h[rs[root]]) + 1,memset (dp[root],0x3f,sizeof (dp[root]));
		for (Int i = h[root];i <= hei[root];++ i)
			chkmin (dp[root][i],dp[ls[root]][i - 1] + dp[rs[root]][i - 1] + (root > x)),
			chkmin (dp[root][i],dp[ls[root]][i - 2] + dp[rs[root]][i - 1] + (root > x)),
			chkmin (dp[root][i],dp[ls[root]][i - 1] + dp[rs[root]][i - 2] + (root > x));
	}
}

signed main(){
	freopen ("avl.in","r",stdin);
	freopen ("avl.out","w",stdout);
	read (n,K);
	for (Int i = 1;i <= n;++ i){
		int p;read (p);
		if (p == -1) rt = i;
		else if (i < p) ls[p] = i;
		else rs[p] = i;
	}
	dfs1 (rt),f[1] = 1,f[2] = 2;
	for (Int i = 3;i <= 25;++ i) f[i] = f[i - 1] + f[i - 2] + 1;
	memset (dp,0x3f,sizeof (dp));
	for (Int u = 0;u <= n;++ u)
		for (Int i = 0;i <= hei[u];++ i) dp[u][i] = f[i];
	for (Int i = 1;i <= n;++ i){
		tot = 0,ins (rt,i);
		if (num + 1 + dp[rt][h[rt]] > K){
			for (Int i = 1;i <= tot;++ i){
				int now = pos[i];
				h[now] = reh[i];for (Int k = 0;k <= hei[now];++ k) dp[now][k] = sta[i][k];
			}
		}
		else ans[i] = 1,num ++;
	}
	for (Int i = 1;i <= n;++ i) putchar (ans[i] + '0');putchar ('\n');
 	return 0;
}