[Advanced/C++/1135] Suggestion

[zzzbf]

一次前序遍历即可，不需要遍历两次

#include<iostream>
#include<vector>
#include<set>

using namespace std;

struct node {
	int val;
	node *left, *right;
};

node *Insert(int val, node *root)
{
	if (root == NULL) {
		root = new node();
		root->val = val;
		root->left = root->right = NULL;
	}
	else if (abs(val) <= abs(root->val))
		root->left = Insert(val, root->left);
	else
		root->right = Insert(val, root->right);
	return root;
}

bool preOrderJudge(node* root, int blackCnt, set<int> &cnt) {
	if (root == NULL) {
		cnt.insert(blackCnt);
		return true;
	}
	
	if(root->val<0 && ((root->right && root->right->val<0)|| 
		(root->left && root->left->val < 0))) return false;
	if (preOrderJudge(root->left, root->val < 0 ? blackCnt : blackCnt + 1, cnt))
		return preOrderJudge(root->right, root->val < 0 ? blackCnt : blackCnt + 1, cnt);
	return false;
}

int main()
{
	int testN, nodesN, val;
	cin >> testN;
	for (int i = 0; i < testN; i++) {
		cin >> nodesN;
		node* root = NULL;
		set<int> cnt;
		for (int j = 0; j < nodesN; j++) {
			cin >> val;
			root = Insert(val, root);
		}
		// 根结点必须为黑结点
		// 红色结点的孩子为黑色结点
		// 所有路径有相同数量的黑色结点

		bool flag = preOrderJudge(root, 0, cnt);
		printf("%s\n", root->val > 0 && flag && cnt.size() == 1 ? "Yes" : "No");
	}
	return 0;
}