Count nodes in tree program


import java.util.ArrayDeque;

class ConstructTree {
    
    static int count(Node root) {
        if(root != null) {
            int leftCount = count(root.left);
            int rightCount = count(root.right);
            return leftCount + rightCount + 1;
        }
        return 0;
    }
    
    static Node construct(int[] preOrder, int[] inOrder) {
        // base check
        if (preOrder.length == 0 || preOrder.length != inOrder.length) {
            return null;
        }
        
        // counters
        int preOrderIndex = 0;
        int inOrderIndex = 0;
        
        ArrayDeque stack = new ArrayDeque<>();
        
        Node root = new Node(preOrder[preOrderIndex]);
        stack.addFirst(root);
        preOrderIndex++;
        
        while(!stack.isEmpty()) {
            Node top = stack.peekFirst();
            if(top.data == inOrder[inOrderIndex]) {
                stack.pollFirst();
                inOrderIndex++;
                
                // if all the elements in inOrder have been visted, we are done
                if (inOrderIndex == inOrder.length) {
                    break;
                }
                
                // Check if there are still some unvisited nodes in the left
                // sub-tree of the top node in the stack
                if (!stack.isEmpty() && stack.peekFirst().data == inOrder[inOrderIndex]) {
                    continue;
                }
                
                // As top node in stack, still has not encontered its counterpart
                // in inOrder, so next element in preOrder must be right child of
                // the removed node
                Node node = new Node(preOrder[preOrderIndex]);
                preOrderIndex++;
                top.right = node;
                stack.addFirst(node);
            } else {
                // Top node in the stack has not encountered its counterpart
                // in inOrder, so next element in preOrder must be left child
                // of this node
                Node node = new Node(preOrder[preOrderIndex]);
                preOrderIndex++;
                top.left = node;
                stack.addFirst(node);
            }
        }
        return root;
    }
    
    static void inOrderTraversal(Node root) {
        if(root == null) {
            return;
        }
        inOrderTraversal(root.left);
        System.out.print(root.data + " ");
        inOrderTraversal(root.right);
    }
    
    static void preOrderTraversal(Node root) {
        if(root == null) {
            return;
        }
        System.out.print(root.data + " ");
        preOrderTraversal(root.left);
        inOrderTraversal(root.right);
    }
    
	public static void main (String[] args) {
	    int[] inOrder = new int[] {4, 2, 5, 8, 1, 6, 3, 9, 7};
        //int[] preOrder = new int[] {1, 2, 4, 5, 8, 3, 6, 7, 9};
        int[] preOrder = new int[] {1, 2, 4, 5, 8, 6, 3, 9, 7};
        
        Node root = construct(preOrder, inOrder);
        
        // traverse the constructed tree
        System.out.print("The inorder traversal is ");
        inOrderTraversal(root);
 
        System.out.print("\nThe preorder traversal is ");
        preOrderTraversal(root);
        
        System.out.print("\nTotal number of nodes : " + count(root));
	}
	
	private static class Node {
	    Node left;
	    int data;
	    Node right;
	    public Node(int data) {
	        this.data = data;
	    }
	}
}

Output
The inorder traversal is 4 2 5 8 1 6 3 9 7 
The preorder traversal is 1 2 4 5 8 6 3 9 7 
Total number of nodes : 9

Comments

Popular posts from this blog

SQL basic interview question

gsutil Vs Storage Transfer Service Vs Transfer Appliance