[PAT顶级1035] 能否染色为红黑树–构建二叉搜索树+位运算树形dp
题目描述
对红黑树的描述和 [PAT甲级1135]判断红黑树(AcWing链接) 完全一致。但是做了两处修改:
- 输入序列由前序遍历改成了后序遍历
- 不给定每个节点已经染了什么颜色,而是让你判断当前这棵树能否被染色成一个红黑树
样例
样例输入 #1
3
9
1 4 5 2 8 15 14 11 7
9
1 4 5 8 7 2 15 14 11
8
6 5 8 7 11 17 15 10
样例输出 #1
Yes
No
Yes
思路
原本的“判断红黑树”一题是边构建二叉搜索树边判断两边的黑色节点路径长度是否一致,那么我们这里就把构建二叉搜索树和判断能否染色这两个过程给分开。
构建二叉搜索树的时候和之前一样,dfs的过程中设置一个当前dfs是在序列的哪个区间,在值域的哪个区间(前提是已经对值域做过离散化并映射到了 $[1, n]$)。跟前序遍历类似,我们可以确定根节点的前一个数就是自己的右子节点,同时根据根节点的值可以判断右子树的规模是多大,从而确定了左子节点以及左子树的规模。规模确定之后,两颗子树的值域区间也确定了。只要出现了根节点的值在值域区间之外就说明并不是一个合法的二叉搜索树。然后我们将每个节点的左右子节点记录下来,就算是完成了树形结构的构造。
然后看是否可能染色为红黑树。“判断红黑树”一题是给定了每个根节点的黑节点路径长度是多少,那么我们在这里就要改成每个根节点可能的黑节点路径长度是多少。那么我们就可以设置对应的dp方案:设 $dp[i][0/1]$ 为 $i$ 节点分别染成黑色和红色时,可能实现的黑节点路径长度是多少。
初始情况下,空节点一定是黑色,所以其 $dp[0][0]=\{0\}, dp[0][1]=\emptyset$ (我们设只有一个节点时的树高为 $0$)
那么每个节点在其左右子节点dp过程结束时开始观察,不难发现如果左右子节点都有至少一个让黑节点路径为 $i$ 的合法染色方案,那么当前节点就一定会有一个让黑节点路径为 $i$ (当前节点是红色) 或者 $i+1$ (当前节点是黑色) 的合法染色方案。
不难发现左右子节点查找共同方案的过程就是位运算中的与运算,而 $i$ 变成 $i+1$ 就是左移运算。再考虑到题目中 $n<30$ , 我们可以把整个方案集合压缩到一个 int
,其二进制的第 $i$ 位表示是否存在黑节点路径为 $i$ 的合法染色方案就行了(如果 $n$ 到了 $1000$ 这个级别,我们也可以用 bitset
) 。
回到dp过程,具体的dp还需要遵守红黑树的规则,父节点是黑色时子节点无所谓,所以就把左黑右黑/左红右黑/左黑右红/左红右红这四个与运算之后的方案在一起做个或运算,然后左移一位。父节点是红节点时那就只有左黑右黑这一个方案,与运算一下就行了。
最后树根还得是黑色的,所以我们就看 $dp[\text{root}][0]$ 是否不是空集(对应到 int
中就是是否为 $0$) 如果不为 $0$ ,那就是有合法染色方案,反之就是没有。
最后注意一下多组数据要对数组做初始化。
代码
#include <stdio.h>
#include <string.h>
#include <algorithm>
int T;
int n;
int post[40], a[40];
int ch[40][2]; // 0 : left child 1 : right child
bool flag; // is BST or not
int dfs_build(int pos_l, int pos_r, int val_l, int val_r)
{
int rt = pos_r, val_rt = post[rt];
if (val_rt < val_l || val_rt > val_r) // illegal BST
{
flag = false;
return 0;
}
ch[rt][0] = ch[rt][1] = 0;
if (val_l < val_rt)
ch[rt][0] = dfs_build(pos_l, pos_r - (val_r - val_rt) - 1, val_l, val_rt - 1);
if (val_r > val_rt)
ch[rt][1] = dfs_build(pos_r - (val_r - val_rt), pos_r - 1, val_rt + 1, val_r);
return rt;
}
int possible_black_height[40][2]; // if node is : 0 black 1 red
void dfs_dp(int u)
{
int lc = ch[u][0], rc = ch[u][1];
if (lc)
dfs_dp(lc);
if (rc)
dfs_dp(rc);
for (int i = 0; i <= 1; ++i)
for (int j = 0; j <= 1; ++j)
possible_black_height[u][0] |= ((possible_black_height[lc][i] & possible_black_height[rc][j]) << 1);
possible_black_height[u][1] = (possible_black_height[lc][0] & possible_black_height[rc][0]);
}
void init()
{
flag = true;
memset(possible_black_height, 0, sizeof(possible_black_height));
possible_black_height[0][0] = 1; // set null node to black, black height is 0
}
int main()
{
scanf("%d", &T);
while (T--)
{
init();
scanf("%d", &n);
for (int i = 1; i <= n; ++i)
scanf("%d", &post[i]), a[i] = post[i];
std::sort(a + 1, a + n + 1);
for (int i = 1; i <= n; ++i)
post[i] = std::lower_bound(a + 1, a + n + 1, post[i]) - a;
int rt = dfs_build(1, n, 1, n);
if (!flag)
{
puts("No");
continue;
}
dfs_dp(rt);
puts(possible_black_height[rt][0] ? "Yes" : "No");
}
}