888. 公平的糖果交換 Fair Candy Swap

題目 https://leetcode-cn.com/problems/fair-candy-swap/

/**
 * Note: The returned array must be malloced, assume caller calls free().
 */
struct Node{
	int key;
	struct Node *left,*right;
	int isRED;
};


struct Node* getValue(struct Node *root,int key){
	if(root == NULL){
		return NULL;
	}
	if(root->key == key){
		return root;
	}
	else if(root->key < key){
		return getValue(root->right,key);
	}else{
		return getValue(root->left,key);
	}
}

struct Node* balance(struct Node* root){
	if(root == NULL){
		return NULL;
	}

	if(root->left != NULL && root->right != NULL && root->left->isRED == 1 && root->right->isRED == 1)
	{
		root->left->isRED = 0;
		root->right->isRED = 0;
		root->isRED = 1;
		return root;
	}

	if(root->left != NULL && root->left->isRED == 1){
		if(root->left->left!=NULL && root->left->left->isRED == 1){
			struct Node *node = root;
			root = node->left;
			struct Node* right = root->right;

			root->left->isRED = 0;
			root->right = node;
			root->right->isRED = 0;
		
			node->left = right;
			return root;
		}
		if(root->left->right!=NULL && root->left->right->isRED == 1){
			struct Node *node = root;
			root = node->left->right;
			struct Node* left = root->left;
			struct Node* right = root->right;

			root->left = node->left;
			root->left->right = left;
			root->left->isRED = 0;
			root->right = node;

			node->left = right;
			return root;
		}
	}
	if(root->right != NULL && root->right->isRED == 1){
		if(root->right->left!=NULL && root->right->left->isRED == 1){
			struct Node *node = root;
			root = root->right->left;
			struct Node* left = root->left;
			struct Node* right = root->right;

			root->left = node;
			root->right = node->right;
			root->right->left=right;
			root->right->isRED = 0;

			node->right=left;
			return root;
		}
		if(root->right->right!=NULL && root->right->right->isRED == 1){
			struct Node *node = root;
			root = node->right;
			struct Node* left = root->left;

			root->left = node;
			root->left->isRED = 0;
			root->right->isRED = 0;

			node->right = left;
			return root;
		}
	}
	return root;
}

struct Node* putValue(struct Node *root,int key){
	if(root == NULL){
		struct Node* node = malloc(sizeof(struct Node));
		memset(node,0,sizeof(node));
		node->key = key;
		node->isRED = 1;
		node->left = NULL;
		node->right = NULL;
		return node;
	}

	if(root->key == key){
		return root;
	}
	else if(root->key < key){
		root->right = putValue(root->right,key);
	}else{
		root->left = putValue(root->left,key);
	}
	root = balance(root);
	return root;
}

void freeTree(struct Node* root){
	if(root == NULL)
		return;
	freeTree(root->left);
	freeTree(root->right);
	free(root);
}

int* fairCandySwap(int* A, int ASize, int* B, int BSize, int* returnSize){
	long Asum = 0,Bsum=0,Halfsum = 0;
	int n,i;
	struct Node *Aroot = NULL,*Broot = NULL;
	struct Node* node = NULL;
	int *C = malloc(sizeof(int)*2);
	for(i=0;i<ASize;i++)
	{
		Asum+=A[i];
		Aroot = putValue(Aroot,A[i]);
	}
	for(i=0;i<BSize;i++)
	{
		Bsum+=B[i];
		Broot = putValue(Broot,B[i]);
	}
	Halfsum = (Asum+Bsum)/2;
	if(ASize < BSize)
	{
		if(Asum > Bsum)
		{
			n = Asum-Halfsum;
			for(i=0;i<ASize;i++)
			{
				node = getValue(Broot,A[i]-n);
				if(node != NULL)
					break;
			}
			C[0] = A[i];
			C[1] = node->key;
		}
		else
		{
			n = Bsum-Halfsum;
			for(i=0;i<ASize;i++)
			{
				node = getValue(Broot,A[i]+n);
				if(node != NULL)
					break;
			}
			C[0] = A[i];
			C[1] = node->key;
		}
	}
	else
	{
		if(Asum > Bsum)
		{
			n = Asum-Halfsum;
			for(i=0;i<BSize;i++)
			{
				node = getValue(Aroot,B[i]+n);
				if(node != NULL)
					break;
			}
			C[0] = node->key;
			C[1] = B[i];
		}
		else
		{
			n = Bsum-Halfsum;
			for(i=0;i<BSize;i++)
			{
				node = getValue(Aroot,B[i]-n);
				if(node != NULL)
					break;
			}
			C[0] = node->key;
			C[1] = B[i];
		}
	}
	freeTree(Aroot);
	freeTree(Broot);
	*returnSize = 2;
	return C;
}