Find the rightmost cousin of a binary tree

Given a node in a binary tree. Find the rightmost cousin of the give node.

Let’s start with an example.

                             1
                          /      \
                         2        3
                      /         /   \
                     4         5     6
                       \           /
                        7         8

Rightmost cousin of 4 is 6, right most cousin of 7 is 8 and so on.

The idea behind solving this problem is to get to the level where the node exists. Clearly we can do this by a level order traversal. Note that, level order traversal is similar to breadth first search i.e. traverse the children of a node before going into deeper. We can use a Queue data structure to save all children of a node in FIFO order in order to traverse in level order. The tricky part here is to find out when a new level begins. We can easily keep track of the level by counting number of nodes to be traversed in current level. Once we reach a node which is the parent of the target node the we know that next level is our desired level and all the nodes in queue are in the same level as target node.

For example, Lets find the cousin of 4. initially the queue has one element at level=1 i.e. the root {1} and count=1. While traversing we remove 1 from the queue. Now count = 0 and a new level=2 began where we will enqueue {2, 3} and new count = size of queue = 2. The next level will begin after two traverse. So, we can identify the beginning of next level by checking that number of traversals in current level is equal to the size of the queue during the end of previous level. For example: During level=2, count=2. So, next level will occur exactly after two remove operation. First 2 will be removed (now count=1) and its children {4} will be added in queue. Also, we know that the next level is our target level as we hit the target node 4. Then 3 will be removed (now count=0) and its children {5, 6} will be added to queue. At this moment new count = size of queue = 3.

Now, we are in the target level and all the nodes in the queue are in the same level hence are cousins and the last node is the right most cousin. Below is a O(n) time O(n) space code by using a LinkedList.

public static class BTNode{
	int key;
	int val;
	BTNode left;
	BTNode right;
	BTNode next;
	
	public BTNode(int key){
		this.key = key;
	}

	@Override
	public String toString() {
		return key+"";
	}
}

public static BTNode rightMostCousin(BTNode root, int targetKey){
	LinkedList<BTNode> q = new LinkedList<BTNode>();
	
	int count = 0;
	q.add(root);	
	count++;
	boolean targetLevel = false;
	
	while(!q.isEmpty())
	{
		BTNode node = q.remove();	
		count--;
		if(node.key == targetKey)
			targetLevel = true;			
		
		if(node.left != null) q.add(node.left);
		if(node.right != null) q.add(node.right);

		if(count == 0){			
			count = q.size();
			if(targetLevel){
				if(node.key != targetKey)
					return node;
				else return null;
			}
		}
	}
	
	return null;
}

Slight variation of this program can lead us to solution to other problems such as:

1. Find all the cousins of a binary tree node
2. Print a binary tree in level order
3. Connect level order sibling
4. Find level of a node, etc.

public static void printLevelOrder(BTNode root){
	Queue<BTNode> queue = new LinkedList<test.BTNode>();
	queue.offer(root);
	
	BTNode node = null;
	int count = 1;
	while(!queue.isEmpty()){
		node = queue.poll();
		count--;
		
		System.out.print(node.key+" ");
		if(node.left != null){
			queue.offer(node.left);
		}
		if(node.right != null){
			queue.offer(node.right);
		}
		
		if(count == 0){
			System.out.println("");
			count = queue.size();
		}
	}
}

public static void connectLevelOrder(BTNode root){
	Queue<BTNode> queue = new LinkedList<test.BTNode>();
	queue.offer(root);
	
	BTNode node = null;
	int count = 1;
	while(!queue.isEmpty()){
		if(node == null){
			node = queue.poll();
			node.next = null;
		}
		else{
			node.next = queue.poll();
			node = node.next;
		}
		count--;
		
		System.out.print(node.key+" ");
		if(node.left != null){
			queue.offer(node.left);
		}
		if(node.right != null){
			queue.offer(node.right);
		}
		
		if(count == 0){
			node = null;
			System.out.println("");
			count = queue.size();
		}
	}
}