Binary Tree utils for (1) getting nodes at a given distance & (2) getting zig zag order traversal of binary tree. Includes two different implementation for both operations.

BinaryTreeUtils.java

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.LinkedList;
import java.util.Map;
import java.util.Stack;
import java.util.Queue;
import java.util.Collections;

/**
 * Binary Tree utils for <br/>
 *  - getting nodes at a given distance <br/>
 *  - getting zig zag order traversal of binary tree <br/>
 *
 * @author ozkansari
 */
public class BinaryTreeUtils {

	/**
	 * BinaryTreeOperations interface instances
	 */
	private BinaryTreeOperations btInstance1, btInstance2 ;
	
	/* ------------------------------------------------------------- */
	/* SOLUTION- 1 */
	/* ------------------------------------------------------------- */

	/**
	 * My BFS-like Implementation
	 * 
	 * @author ozkansari
	 *
	 */
	public class BinaryTreeOperationsBFSImpl implements BinaryTreeOperations {

		/**
		 * BFS-like implementation
		 *
		 * Total Time complexity: O(2n)
		 * Total Space complexity: O(3n) -> size(bfsQueue) + size(resultList) + size(levelNodes)
		 */
		@Override
		public List<Node> getNodesAtDistance(Node rootNode, int distance){
			List<Node> resultList = new ArrayList<Node>();

			Queue<Node> bfsQueue = new LinkedList();
			bfsQueue.add(rootNode);

			int currentLevel = 0;
			boolean hasChildNodes;
			while(!bfsQueue.isEmpty()) {

				hasChildNodes = false;

				// O(n)
				// Remove all first
				List<Node> currentLevelNodes = new ArrayList<Node>();
				while(!bfsQueue.isEmpty()) {
					currentLevelNodes.add(bfsQueue.remove());
				}

				// O(n)
				for (Node n : currentLevelNodes) {
					Node leftNode = n.getLeft();
					Node rightNode = n.getRight();

					if(leftNode!=null) {
						bfsQueue.add(leftNode);
						hasChildNodes = true;
					}

					if(rightNode!=null) {
						bfsQueue.add(rightNode);
						hasChildNodes = true;
					}
					
				}

				// check return condition
				if(currentLevel==distance) {
					resultList.addAll(currentLevelNodes);
					break;
				}

				// If has more childs, increment the level
				if(hasChildNodes) {
					currentLevel++;
				}

			}
			return resultList;
		}

		/**
		 * Total Time complexity: O(2n) -> n+n+(n/8)
		 * Total Space complexity: O(3n) -> size(bfsQueue) + size(resultList) + size(levelNodes) 
		 */
		@Override
		public List<Node> getZigZagOrder(Node rootNode) {
			Queue<Node> bfsQueue = new LinkedList<Node>();

			List<Node> resultList = new ArrayList<Node>();
			bfsQueue.add(rootNode);
			// resultList.add(rootNode);

			int currentLevel = 1;
			while(!bfsQueue.isEmpty()) {

				// O(n)
				// Remove all first
				List<Node> levelNodes = new ArrayList<Node>();
				while(!bfsQueue.isEmpty()) {
					levelNodes.add(bfsQueue.remove());
				}

				// O(n)
				boolean hasChildNode = false;
				for (Node currentNode : levelNodes) {
					Node firstNode = currentNode.getLeft();
					Node secondNode = currentNode.getRight();

					if(firstNode!=null) {
						bfsQueue.add( firstNode );
						hasChildNode=true;
					}
						
					if(secondNode!=null) {
						bfsQueue.add( secondNode );
						hasChildNode=true;
					}

				}

				// O(n/8)
				if(currentLevel%2==0) {
					Collections.reverse(levelNodes);
				}
				resultList.addAll(levelNodes);

				if(hasChildNode) {
					currentLevel++;
				}

			}

			return resultList;
		}
	}

	/* ------------------------------------------------------------- */
	/* SOLUTION- 2 */
	/* ------------------------------------------------------------- */

	/**
	 * My BinaryTreeOperations Implementation
	 * 
	 * @author ozkansari
	 *
	 */
	public class BinaryTreeOperationsStackImpl implements BinaryTreeOperations {

		/**
		 * Map of stacks to represent tree levels
		 */
		Map<Integer,Stack<Node>> levels=new HashMap<Integer,Stack<Node>>();
		
		/*
		 * Recursive solution for getting nodes at a given distance
		 *
		 * Total Time complexity: O(logn)
		 * Total Space complexity: O(n) -> size(resultList)
		 */
		@Override
		public List<Node> getNodesAtDistance(Node rootNode, int distance) {

			// recursive soluiton 
			List<Node> resultList = new ArrayList<Node>();

			if(rootNode==null) {
				return resultList;
			} else if(distance==0) {
				resultList.add(rootNode);
			} else {
				resultList.addAll( getNodesAtDistance(rootNode.getLeft(), distance-1) );
				resultList.addAll( getNodesAtDistance(rootNode.getRight(), distance-1) );
			}
			// System.out.println(rootNode + " - r:" + distance + " , size: " + resultList.size());
			return resultList;

		}

		/*
		 * Total Time complexity: O(n) -> n+n
		 * Total Space complexity: O(n) -> size(levels)
		 */
		@Override
		public List<Node> getZigZagOrder(Node rootNode) {
			
			// Initialize next level
			initLevels(rootNode, 1);
			
			// Find result list
			List<Node> resultList = new ArrayList<Node>();

			for (int i=1; ; i++) {
				
				Stack<Node> levelStack = levels.get(i);
				if(levelStack==null) break;
				
				while(!levelStack.empty()) {
					Node currentNode = null;
					
					if(i%2==0) { // right side
						currentNode = levelStack.pop();
					} else {
						currentNode = levelStack.firstElement();
						levelStack.remove(currentNode);
					}
					
					resultList.add(currentNode);
				}

			}
			
			return resultList;
		}

		/**
		 * Helper recursive method
		 *
		 * @param currentNode
		 * @param currentLevel
		 */
		private void initLevels(Node currentNode, Integer levelIndex) {

			if(levelIndex<=0) {
				levelIndex=1;
			}
			
			// Initialize first root level if required
			if(levelIndex==1){
				Stack<Node> rootLevel = new Stack<Node>();
				rootLevel.add(currentNode);
				levels.put(1,rootLevel);
				levelIndex=2;
			}
			
			Node leftNode = currentNode.getLeft();
			Node rightNode = currentNode.getRight();

			if(leftNode==null && rightNode==null) {
				return;
			} 
			
			Stack<Node> currentLevel = levels.get(levelIndex);
			if(currentLevel==null) {
				currentLevel = new Stack<Node>();
				levels.put(levelIndex,currentLevel);
			}

			if(leftNode!=null) {
				currentLevel.push(leftNode);
				if(leftNode.getLeft()!=null || leftNode.getRight()!=null) {
					initLevels(leftNode,levelIndex+1);
				} 
			}
			
			if(rightNode!=null) {
				currentLevel.push(rightNode);
				if(rightNode.getLeft()!=null || rightNode.getRight()!=null) {
					initLevels(rightNode,levelIndex+1);
				}
			}

		}

	}
	
	/**
	 * 
	 * @return BinaryTreeOperations BFS implementation instance
	 */
	public BinaryTreeOperations getBinaryTreeOperationsBFSImplInstance(){
		if(btInstance1==null) {
			btInstance1 = new BinaryTreeOperationsBFSImpl();
		}
		return btInstance1;
	}

	/**
	 * 
	 * @return BinaryTreeOperations stack implementation instance
	 */
	public BinaryTreeOperations getBinaryTreeOperationsStackImplInstance(){
		if(btInstance2==null) {
			btInstance2 = new BinaryTreeOperationsStackImpl();
		}
		return btInstance2;
	}
	

	/* ------------------------------------------------------------- */
	/* DATA DEFINITION(S) */
	/* ------------------------------------------------------------- */

	/**
	 * Tree node
	 * 
	 * @author an amazonian
	 */
	public class Node {
		
		private int value;
		private Node left;
		private Node right;

		public Node(int value, Node left, Node right) {
			this.value = value;
			this.left = left;
			this.right = right;
		}

		public int getValue() {
			return this.value;
		}

		public Node getLeft() {
			return this.left;
		}

		public Node getRight() {
			return this.right;
		}
		
		protected void setLeft(Node leftNode) {
			this.left = leftNode;
		}

		protected void setRight(Node rightNode) {
			this.right = rightNode;
		}

		@Override
		public String toString() {
			return String.valueOf( getValue() );
		}
	}

	/**
	 * BinaryTreeOperations interface
	 * 
	 * @author an amazonian
	 *
	 */
	public interface BinaryTreeOperations {

		/**
		 *
		 * Print a binary tree in zig zag way... that is:
		 * ......a.........
		 * ....b....c.......
		 * ..d...e...f...g..
		 * .h.i.j.k.l.m.n.o.
		 * 
		 * should be printed as a-c-b-d-e-f-g-o-n-m-l-k-j-i-h
		 * 
		 * what data structure will u use for that?
		 *
		 * @param rootNode
		 *
		 */
		public List<Node> getZigZagOrder(Node rootNode);

		/**
		 * Given a root of a tree, find the nodes at d distance from a certain node in a Binary Tree
		 *
		 * @param rootNode
		 * @param distance
		 *
		 */
		public List<Node> getNodesAtDistance(Node rootNode, int distance);
		
	}


	/* ------------------------------------------------------------- */
	/* TEST METHOD(S) */
	/* ------------------------------------------------------------- */

	/**
	 * Main method to test the BinaryTreeOperations coding sample. 
	 * I actually prefer to test my code using JUnit but the result is requested to be in one file.
	 * 
	 * @author ozkansari
	 * 
	 */
	public String testZigZagOrderTraversal(int nodeCount, BinaryTreeOperations btImpl) {
		
		Node rootNode = generateDummyBinaryTree(nodeCount);
		List<Node> zigZagOrderedList = btImpl.getZigZagOrder(rootNode);
		
		StringBuffer sbf = new StringBuffer();
		for (Node node : zigZagOrderedList) {
			sbf.append(node + " ");
		}
		System.out.println(sbf.toString());
		return sbf.toString();
		
	}

	/**
	 * Main method to test the BinaryTreeOperations coding sample. 
	 * I actually prefer to test my code using JUnit but the result is requested to be in one file.
	 * 
	 * @author ozkansari
	 * 
	 */
	public String testFindNodesAtADistance(int nodeCount, int distance, BinaryTreeOperations btImpl) {
		
		Node rootNode = generateDummyBinaryTree(nodeCount);
		List<Node> nodesAtDistanceList = btImpl.getNodesAtDistance(rootNode,distance);
		
		StringBuffer sbf = new StringBuffer();
		for (Node node : nodesAtDistanceList) {
			sbf.append(node + " ");
		}
		System.out.println(sbf.toString());
		return sbf.toString();
		
	}

	/**
	 * Generates a dummy binary tree and returns its root node
	 * 
	 */
	private Node generateDummyBinaryTree(int nodeCount){
		List<Node> sampleData = new ArrayList<Node>();
		for(int i=1;i<=nodeCount;i++) {
			Node currentNode = new Node(i,null,null);;
			sampleData.add(currentNode);
		}
		int j=1;

		for (Node node : sampleData) {
			int index = j*2;

			Node rightNode = index+1<=sampleData.size() ? sampleData.get(index) : null;
			Node leftNode = index<=sampleData.size() ? sampleData.get(index-1) : null;

			node.setRight(rightNode);
			node.setLeft(leftNode);

			// System.out.println(node + "->" + leftNode + " " + rightNode);
			j++;
		}

		return sampleData.get(0);
	}

	/**
	 * Main method to test my sample code.
	 * 
	 * @param args
	 */
	public static void main(String[] args) {
		BinaryTreeUtils bt = new BinaryTreeUtils();
		BinaryTreeOperations btImpl1 = bt.getBinaryTreeOperationsBFSImplInstance();
		BinaryTreeOperations btImpl2 = bt.getBinaryTreeOperationsStackImplInstance();

		// Test finding nodes at a distance Using BFS Implementation
		System.out.println("______________________________________");
		System.out.println("Test finding nodes at a distance Using BFS Implementation");
		System.out.println("______________________________________");
		System.out.println( bt.testFindNodesAtADistance(1,0,btImpl1).equals("1 ") );
		System.out.println( bt.testFindNodesAtADistance(4,1,btImpl1).equals("2 3 ") );
		System.out.println( bt.testFindNodesAtADistance(12,2,btImpl1).equals("4 5 6 7 ") );
		System.out.println( bt.testFindNodesAtADistance(12,3,btImpl1).equals("8 9 10 11 12 ") );
		System.out.println( bt.testFindNodesAtADistance(12,4,btImpl1).equals("") );

		// Test finding nodes at a distance Using Stack Implementation
		System.out.println("______________________________________");
		System.out.println("Test finding nodes at a distance Using Stack Implementation");
		System.out.println("______________________________________");
		System.out.println( bt.testFindNodesAtADistance(1,0,btImpl2).equals("1 ") );
		System.out.println( bt.testFindNodesAtADistance(4,1,btImpl2).equals("2 3 ") );
		System.out.println( bt.testFindNodesAtADistance(12,2,btImpl2).equals("4 5 6 7 ") );
		System.out.println( bt.testFindNodesAtADistance(12,3,btImpl2).equals("8 9 10 11 12 ") );
		System.out.println( bt.testFindNodesAtADistance(12,4,btImpl2).equals("") );

		// Test ZigZag Order Traversal Using BFS Implementation
		System.out.println("______________________________________");
		System.out.println("Testing ZigZag Order Traversal Using BFS Implementation");
		System.out.println("______________________________________");
		System.out.println( bt.testZigZagOrderTraversal(1,btImpl1).equals("1 ") );
		System.out.println( bt.testZigZagOrderTraversal(2,btImpl1).equals("1 2 ") );
		System.out.println( bt.testZigZagOrderTraversal(4,btImpl1).equals("1 3 2 4 ") );
		System.out.println( bt.testZigZagOrderTraversal(12,btImpl1).equals("1 3 2 4 5 6 7 12 11 10 9 8 ") );
		System.out.println( bt.testZigZagOrderTraversal(16,btImpl1).equals("1 3 2 4 5 6 7 15 14 13 12 11 10 9 8 16 ") );
		System.out.println( bt.testZigZagOrderTraversal(24,btImpl1).equals("1 3 2 4 5 6 7 15 14 13 12 11 10 9 8 16 17 18 19 20 21 22 23 24 ") );

		// Test ZigZag Order Traversal Using Stack Implementation
		System.out.println("______________________________________");
		System.out.println("Testing ZigZag Order Traversal Using Stack Implementation");
		System.out.println("______________________________________");
		System.out.println( bt.testZigZagOrderTraversal(1,btImpl2).equals("1 ") );
		System.out.println( bt.testZigZagOrderTraversal(2,btImpl2).equals("1 2 ") );
		System.out.println( bt.testZigZagOrderTraversal(4,btImpl2).equals("1 3 2 4 ") );
		System.out.println( bt.testZigZagOrderTraversal(12,btImpl2).equals("1 3 2 4 5 6 7 12 11 10 9 8 ") );
		System.out.println( bt.testZigZagOrderTraversal(16,btImpl2).equals("1 3 2 4 5 6 7 15 14 13 12 11 10 9 8 16 ") );
		System.out.println( bt.testZigZagOrderTraversal(24,btImpl2).equals("1 3 2 4 5 6 7 15 14 13 12 11 10 9 8 16 17 18 19 20 21 22 23 24 ") );
	}

}