颜色平衡树C++(后序遍历+一点DP思想）

解题思路:

小白第一次做这种类型的题，u1s1暴力真的香，虽然过不了最后9分但真的容易太多了，过最后一段数据需要用到一点DP思想，将dp[i]视为第i个节点的子树中颜色平衡树的数量，可知叶子节点其本身一定为颜色平衡树，其余节点的颜色平衡树的数量等于各个子节点颜色平衡树的数量之和，再判断以该节点为起点时的搜索树是否为颜色平衡树，若该节点对应的数为颜色平衡树则加1否则不加。

 状态转移方程为：dp[父]=∑dp[子]+{  0,   judge(父)== false;

                                                         1,   judge(父)== true;

其中judge为判断父节点为起始节点的树是否为颜色平衡树。

该题考察的应该是树上启发式合并，应该是优于这种遍历做法的，建议可以去看一下树上启发式合并的解法。



参考代码:

#include <iostream>
#include <vector>
#include <map>
using namespace std;

/*节点结构体*/
//color为颜色值，dp为当前节点的子树中颜色平衡树的数量，count用于存储各颜色数量，child为子节点
struct TreeNode {
    int color;
    map<int,int> count;
    vector<TreeNode*> child;
    TreeNode(int c) :color(c) {};
};

//后序遍历树
int postorder(TreeNode* root) {

    //当遇到叶子节点时返回
    if (root->child.empty()) {
        root->count[root->color]++;
        return 1;
    }
    int sum = 0;

    //遍历当前节点的子节点，将子节点的颜色平衡树累加，并将子节点的颜色计数加入当前节点的颜色计数
    for (vector<TreeNode*>::iterator i = root->child.begin(); i != root->child.end(); i++) {
        sum += postorder(*i);
        for (map<int,int>::iterator j = (*i)->count.begin(); j != (*i)->count.end(); j++)
            root->count[j->first] += j->second;
    }

    //将当前节点的颜色加入计数
    root->count[root->color]++;

    //遍历颜色计数map，为颜色平衡树时令当前节点的dp+1

    int record = root->count.begin()->second;
    int c = 1;
    for (map<int,int>::iterator it = root->count.begin(); it != root->count.end(); it++)
    {
        if (it->second != record) {
            c = 0;
            break;
        }
    }

    return sum + c;
}

int main()
{
    int n;
    cin.sync_with_stdio(false);
    cin.tie(0);
    cin >> n;
    vector<TreeNode*> tree;
    int tmp, fat;
    cin >> tmp >> fat;
    TreeNode* root = new TreeNode(tmp);
    tree.push_back(root);
    for (int i = 1; i < n; i++) {
        int tmp, fat;
        cin >> tmp >> fat;
        TreeNode* node = new TreeNode(tmp);
        tree[fat - 1]->child.push_back(node);
        tree.push_back(node);
    }
    int sum = postorder(root);
    cout << sum;
    return 0;
}