2022省选前联考 AVL树/平衡树

题目描述

pks 得到了一棵 \(N\) 个节点，权值为 \(1\sim N\) 的 \(AVL\) 树，他觉得这棵树太大了，于是他想要删掉一些节点使得最后剩下的树恰好有 \(K\) 个节点。如果 pks 删掉了一个节点，那么以这个节点为根的整棵子树都会被删掉。最后剩下的树必须依旧是一棵 \(AVL\) 树。
pks 希望，留下的 \(K\) 个节点的中序遍历的字典序最小。他希望你能帮他找到这个方案，作为报答，他将会把自己的财富分一半给你。

第一行两个整数 \(N,K\)，表示节点数量和要保留的节点数量。
接下来 \(N\) 行，每行一个数字 \(p_i\)，表示权值为 \(i\) 的节点的父亲的权值，如果是 \(-1\) ，表示这个点是根节点。

一行，\(N\) 个字符，如果权值为 \(i\) 的节点留下来，则第 \(i\) 个字符为 \(1\)，否则为 \(0\) 。

\(N \leq 5 \times 10 ^ 5\)

思路分析

考场只会写暴力
既然要求字典序最小的答案，一种很容易想到的贪心是枚举 \(1 \sim N\) 的节点，如果可以选择就把它加入最后的 \(AVL\) 树中去。但是我们应该如何判断呢？

这里给出一种思路，我们先假设当前节点已经选入 \(AVL\) 树中，再来求出如果它真的被我们选中的话一共至少要选多少个节点（\(AVL\) 树的高度差绝对值不超过一）。
我们可以用 \(dp\) 来解决这个问题。定义 \(f_{i,j}\) 表示以 \(i\) 为根节点，\(AVL\) 树的高度为 \(j\) 最少要选的个数。注意：\(j\) 的大小不应该大于 \(i\) 为根子树的高度大小

我们再从 \(AVL\) 树的定义出发，分别转移 ：

• 左子树，右子树高度相等
• \(左子树高度 - 右子树高度 = 1\)
• \(右子树高度 - 左子树高度 = 1\)

这三种情况。

我们每次判断时都要从当前节点向上跳直到根节点，并且一路标号。
因为如果一个节点选择了，那么它的爸爸，爷爷，曾爷爷 ……一定必选。
每当我们路过一个节点时，我们都要临时更新当前节点的 \(f_{now,j}\) 数组。当选择失败时再还原即可。

#include <bits/stdc++.h>

#define file(a) freopen(a".in", "r", stdin), freopen(a".out", "w", stdout)

#define Enter putchar('\n')
#define quad putchar(' ')

namespace IO {
template <class T> inline void read(T &a);
template <class T, class ...rest> inline void read(T &a, rest &...x);
template <class T> inline void write(T x);
template <class T, class ...rest> inline void write(T x, rest ...a);
}

#define N 500005
#define int long long 

int n, k, root, fa[N], rc[N], lc[N];
int f[N][30], dep[N], g[N][30], flag[N];

inline void update(int now) {
  for (int i = 2; i <= dep[now]; i++) {
    int min1, min2;
    min1 = f[lc[now]][i - 1] + std::min(f[rc[now]][i - 1], f[rc[now]][i - 2]) + 1;
    min2 = f[lc[now]][i - 2] + f[rc[now]][i - 1] + 1;
    f[now][i] = std::min(min1, min2);
  }
  f[now][1] = f[now][0] = 0x3f3f3f3f;
}

inline void dfs(int now) {
  if (now == 0) return ;
  dfs(lc[now]); dfs(rc[now]);
  dep[now] = std::max(dep[lc[now]], dep[rc[now]]) + 1;
  update(now);
  f[now][0] = 0, f[now][1] = 1;
}

inline bool check(int now) {
  if (flag[now]) return true;
  flag[now] ++;
  for (int i = 0; i <= dep[now]; i++) g[now][i] = f[now][i];
  f[now][0] = 0x3f3f3f3f;
  while (now != root) {
    now = fa[now];
    for (int i = 0; i <= dep[now]; i++) g[now][i] = f[now][i];
    update(now); flag[now] ++;
  }
  int minn = 0x3f3f3f3f;
  for (int i = 1; i <= dep[root]; i++)
    minn = std::min(minn, f[root][i]);
  return minn <= k;
}

inline void CtrlZ(int now) {
  flag[now] --;
  for (int i = 0; i <= dep[now]; i++) f[now][i] = g[now][i];
  while (now != root) {
    now = fa[now];
    for (int i = 0; i <= dep[now]; i++) f[now][i] = g[now][i];
    flag[now] --;
  }
}

signed main(void) {
//	file("5071");
  IO::read(n, k);
  for (int i = 1, father; i <= n; i++) {
    IO::read(father);
    if (father == -1) root = i;
    fa[i] = father;
    if (lc[father] == 0) lc[father] = i;
    else rc[father] = i;
  }
  memset(f, 0x3f, sizeof(f));
  for (int i = 0; i <= n; i++) f[i][0] = 0;
  dfs(root);
  for (int i = 1; i <= n; i++) {
    if (check(i)) printf("1");
    else printf("0"), CtrlZ(i);
  }
}

namespace IO {
template <class T> inline void read(T &a) {
  T s = 0, t = 1;
  char c = getchar();
  while ((c < '0' || c > '9') && c != '-')
    c = getchar();
  if (c == '-')
    c = getchar(), t = -1;
  while (c >= '0' && c <= '9')
    s = (s << 1) + (s << 3) + (c ^ 48), c = getchar();
  a = s * t;
}
template <class T, class ...rest> inline void read(T &a, rest &...x) {
  read(a); read(x...);
}

template <class T> inline void write(T x) {
  if (x == 0) putchar('0');
  if (x < 0) putchar('-'), x = -x;
  int top = 0, sta[50] = {0};
  while (x)
    sta[++top] = x % 10, x /= 10;
  while (top)
   putchar(sta[top] + '0'), top --;
   return ;
}
template <class T, class ...rest> inline void write(T x, rest ...a) {
  write(x); quad; write(a...);
}
}