PAT 顶级 1035 - Color the tree / 树的遍历 + 树形DP + 哈希表

原题链接

 https://pintia.cn/problem-sets/994805148990160896/exam/problems/1697152161245261826?type=7&page=0

题目大意

给定一棵有 $N$ 个结点的二叉搜索树的后序遍历，然后给其中所有的结点染色，问是否能染成一棵合法的红黑树。

一棵合法的红黑树应该满足以下几个要求：
（1）每个结点的颜色不是红的就是黑的；
（2）根是黑的；
（3）每个外部叶子结点都是黑色的（外部叶子结点是空结点，比如对于树内部的叶子节点，认为其还有一左一右两个空的节点，这两个空节点就是外部叶子结点、是黑色的）；
（4）如果某个结点是红色的，那么它的所有孩子都是黑色的；
（5）对于某个结点来说，从它开始向下去往任意一个外部叶子结点的路径上，经过的黑色结点数量（也称之为黑高）都相同。

输入格式

输入第一行为一个数 $k$ $(1\le k\le 10)$，表示测试用例个数。
对于每一组测试用例，输入第一行为一个数 $N$ $(1\le N\le 30)$，表示该二叉搜索树的结点个数。输入第二行为该树的后序遍历序列。

输出格式

如果能染成一棵合法的红黑树，则输出“Yes”，否则输出“No”。

输入样例

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

输出样例

Yes
No
Yes

题解

本次PAT顶级考试的两道35分题都不难，思维难度和代码难度都不高，对于DP这个知识点比较熟的同学来说满分应该是手拿把攥的事。通过二叉搜索树的后序遍历生成原树过于简单，就不再叙述了。

本题是一道树形DP问题，但是结合了红黑树的知识点，本题就要从红黑树的性质5入手开始解决。

树形DP的套路是后续遍历，我们从底向上递归地处理这棵树。处理什么呢？处理树的黑高。红黑树的性质5描述，每个结点到外部叶子结点路径上黑色结点的数量都应该相同。也就是说，如果一棵红黑树能够成立，对于任意一个结点，至少它的左右子树的黑高应该相等，这样就将原问题拆分成了两个相同的子问题。

由于每个结点孩子的颜色和黑高都是待定的，因此我们用两个哈希表分别记录每个结点分别染成红色和黑色可能的黑高即可。比如一个内部叶子节点，如果它染红黑高就是$1$，如果它染黑黑高就是$2$，所以将$1$和$2$都记录到哈希表里。

我们假设在某一时刻递归到了结点$root$，我们来分析一下给$root$染色，一共有几种可能的情况不与红黑树的条件冲突：
Case 1：$root$染黑，此时两个孩子结点随便染什么颜色都可以，只要能保证两边的黑高相等就行，所以有四种情况：左黑右黑、左黑右红、左红右黑、左红右红，两两匹配哈希表，如果发现其中有相等元素，$root$中就可以加入一个元素，它的黑高就是那个相等的高度+1。需要注意的细节是，如果孩子是空结点，其实是默认它是一个黑高为1的外部叶子节点，所以这个时候要额外记录一下黑高可以为1。
Case 2：$root$染红，此时根据性质4，两个孩子结点必须染黑，匹配一次两个孩子染黑的哈希表即可。

如果最后有一个黑高是合理的，会被记录在根节点的两个哈希表中。条件2虽然说根必须是黑的，但是根是红的也不会影响任何的黑高，直接将其变成黑的也毫无影响，所以不需要额外去考虑这个条件，如果根节点的任何一个哈希表中存在元素，就证明可以得到合法的红黑树，输出“Yes”，否则就输出“No”。

相关链接

树形DP是有比较经典的套路的：定义状态、后序遍历、状态机/状态转移，不熟悉的同学可以做一下这几道题练习一下。

没有上司的舞会 https://www.luogu.com.cn/problem/P1352
打家劫舍III https://leetcode.cn/problems/house-robber-iii/
最小化旅行的价格总和 https://leetcode.cn/problems/minimize-the-total-price-of-the-trips/
完美树 https://pintia.cn/problem-sets/994805046380707840/exam/problems/1649748772845703169?type=7&page=1

AC代码

#include <bits/stdc++.h>
using namespace std;
struct TreeNode{
    map<int,int> redheight,blackheight;
    int val;
    TreeNode* left, *right;
};
TreeNode* build(int *post,int n)
{
    if(n == 0) return nullptr;
    int i = 0;
    while(post[i]<post[n-1]) i++;
    TreeNode* root = new TreeNode;
    root->val = post[n-1];
    root->left = build(post,i);
    root->right= build(post+i,n-1-i);
    return root;
}
void traversal(TreeNode* root)
{
    if(!root) return;
    traversal(root->left);
    traversal(root->right);
    map<int,int> blackleft,blackright,redright,redleft;
    if(!root->left){
        blackleft[1] = 1;
    } else {
        blackleft = root->left->blackheight;
        redleft = root->left->redheight;
    }
    if(!root->right){
        blackright[1] = 1;
    } else {
        blackright = root->right->blackheight;
        redright = root->right->redheight;
    }
    //如果这个结点涂黑，左右两边可涂任意颜色，但是要黑高一致
    for(auto k:blackleft){
        for(auto h:blackright){
            if(k.first == h.first) root->blackheight[k.first+1] ++;
        }
        for(auto h:redright){
            if(k.first == h.first) root->blackheight[k.first+1] ++;
        }
    }
    for(auto k:redleft){
        for(auto h:blackright){
            if(k.first == h.first) root->blackheight[k.first+1] ++;
        }
        for(auto h:redright){
            if(k.first == h.first) root->blackheight[k.first+1] ++;
        }
    }    
    //如果这个结点涂红，左右两边必须全是黑的
    for(auto k:blackleft){
        for(auto h:blackright){
            if(k.first== h.first) root->redheight[k.first]++;
        }
    }
}
int main()
{
    int t;
    cin >> t;
    while(t--){
        int n;
        cin >> n;
        int post[31];
        for(int i=0;i<n;i++) cin >> post[i];
        TreeNode* root = build(post,n);
        traversal(root);
        if(root->blackheight.size() || root->redheight.size()){
            cout << "Yes" << endl;
        } else cout << "No" << endl;
    }
}