Count pairs from two binary search trees whose sum is equal to a given value X

Article Creation Date : 22-Jun-2021 03:39:45 PM

Given two BSTs containing n1 and n2 distinct nodes respectively. Given a value x. The problem is to count all pairs from both the BSTs whose sum is equal to x.We can implement this by selecting one node value 'a' from BST 1 and search (x-a) value in BST 2, if found increment the count.

Program:

 //C++ program to count pairs from two binary search trees whose sum is equal to a given value x
#include<bits/stdc++.h>
using namespace std;

class TNode
{
	public:
	int data;
	TNode* left;
	TNode* right;
};

TNode* newNode(int data);

TNode* sortedArrayToBST(int arr[],int start, int end)
{
	if (start > end)
	return NULL;
	int mid = (start + end)/2;
	TNode *root = newNode(arr[mid]);

	root->left = sortedArrayToBST(arr, start,mid - 1);
	root->right = sortedArrayToBST(arr, mid + 1, end);

	return root;
}

TNode* newNode(int data)
{
	TNode* node = new TNode();
	node->data = data;
	node->left = NULL;
	node->right = NULL;

	return node;
}

int countPairs(TNode* root1, TNode* root2, int x)
{
	
	if (root1 == NULL || root2 == NULL)
		return 0;

	stack<TNode*> st1, st2;
	TNode* top1, *top2;

	int count = 0;

	while (1) {

		while (root1 != NULL) {
			st1.push(root1);
			root1 = root1->left;
		}
	                 while (root2 != NULL) {
			st2.push(root2);
			root2 = root2->right;
		}

		if (st1.empty() || st2.empty())
			break;

		top1 = st1.top();
		top2 = st2.top();

		if ((top1->data + top2->data) == x) {
			
			count++;

			st1.pop();
			st2.pop();

			root1 = top1->right;
			root2 = top2->left;
		}

		else if ((top1->data + top2->data) < x) {
			st1.pop();
			root1 = top1->right;
		}

		else {
			st2.pop();
			root2 = top2->left;
		}
	}

	return count;
}

void preOrder(TNode* node)
{
	if (node == NULL)
		return;
	cout << node->data << " ";
	preOrder(node->left);
	preOrder(node->right);
}

int main()
{
	int arr[] = {1, 2, 3, 4, 5, 6, 7};
	int n = sizeof(arr) / sizeof(arr[0]);

	int arr1[] = {3,6,8,10,11,15,18};
	int n1 = sizeof(arr1) / sizeof(arr1[0]);

	/* Convert List to BST */
	TNode *root1 = sortedArrayToBST(arr, 0, n-1);
	cout << "PreOrder Traversal of constructed BST \n";
	preOrder(root1);

	TNode *root2 = sortedArrayToBST(arr1, 0, n1-1);
	cout << "\nPreOrder Traversal of constructed BST \n";
	preOrder(root2);

	int x = 16;
	cout << "Pairs found = "
		<< countPairs(root1, root2, x);

	return 0;
}

Output:

PreOrder Traversal of constructed BST1                                                        
4 2 1 3 6 5 7                                                                                
PreOrder Traversal of constructed BST2                                                        
10 6 3 8 15 11 18 
Pairs found = 3