[GeeksForGeeks] Check if removing an edge creates two trees of equal size

Given a Binary Tree, find if there exist edge whose removal creates two trees of equal size.

Examples:

Input : root of following tree
           5
         /   
       1      6    
      /      /  
     3      7    4
Output : true
Removing edge 5-6 creates two trees of equal size


Input : root of following tree
           5
         /   
       1      6    
            /  
           7    4
         /      
        3    2    8
Output : false
There is no edge whose removal creates two trees
of equal size.

To find whether such an edge exists, we need to find if there is at least one node that meets the following property.

either this node's left subtree or right subtree has exactly half of the total number of nodes in the given binary tree.

Solution 1. Top down, O(n^2) runtime

1.  Count the total number of nodes in given binary tree in O(n) time.

2.  from top to bottom, check each node to see if it meets the required property. Each check takes O(n) time, so to check

all nodes, it takes O(n^2) time.

The bottleneck in this algorithm is the O(n) of checking required property operation. For any node, this algorithm recounts 

the total number of nodes of its left and right subtree, FROM SCRATCH. For example, for node A and its left child B,

to check if A meets the property, we need to count how many nodes there are in A's left subtree; 

to check if B meets the property, we need to count how many nodes there are in B's left/right subtree.

B's subtree nodes are counted twice already: once for when counting A's left subtree; once for when counting B's subtrees. 

To avoid this redundant work, we should count bottom up and use memoization to get the number of nodes of given node's subtree in O(1) time.

This is demonstrated in Solution 2.

public class RemoveEdgeToEqualSize {
    public boolean isAnySuchEdgeInefficient(TreeNode root) {
        int totalNodes = countTotalNodes(root);        
        return checkAllNodes(root, totalNodes);
    }
    private int countTotalNodes(TreeNode node) {
        if(node == null) {
            return 0;
        }
        int count = 1;
        count += countTotalNodes(node.left);
        count += countTotalNodes(node.right);
        return count;
    }
    private boolean checkOneNode(TreeNode node, int totalNodes) {
        if(node == null) {
            return false;
        }
        int leftSubtreeNodes = countTotalNodes(node.left);
        if(leftSubtreeNodes * 2 == totalNodes) {
            return true;
        }
        int rightSubtreeNodes = countTotalNodes(node.right);
        if(rightSubtreeNodes * 2 == totalNodes) {
            return true;
        }
        return false;
    }
    private boolean checkAllNodes(TreeNode node, int totalNodes) {
        if(node == null) {
            return false;
        }
        if(checkOneNode(node, totalNodes)) {
            return true;
        }
        return checkAllNodes(node.left, totalNodes) || checkAllNodes(node.right, totalNodes);
    }
}

Solution 2. Bottom up, O(n) runtime

1. Count the total number of nodes in given binary tree in O(n) time.

2.  From bottom to top,  for each node A, treat it as a root node and count the total number of nodes in this subtree recursively.

3.  For each node A, check if there exists an edge that halves the given input tree in this subtree of root A. If any of the 3 is true,

then there exists such an edge in subtree of root A.

　　a. subtree of root A has half of all nodes of the given input tree.

　　b. there exists such an edge in A's left subtree.

　　c.  there exists such an edge in A's right subtree.

For each node A, it takes O(1) time to calcuate its total number of nodes in subtree of root A. It also takes O(1) time to deterimine

if there exists such an edge in subtree of root A. So it takes O(n) time to check all nodes in given binary tree.

class ReturnInfo {
    int totalCnt;
    boolean isAnySuchEdge;
    ReturnInfo(int t, boolean b) {
        totalCnt = t;
        isAnySuchEdge = b;
    }
}
public class RemoveEdgeToEqualSize {
    public boolean isAnySuchEdgeEfficient(TreeNode root) {
        int totalNodes = countTotalNodes(root);
        return recurHelper(root, totalNodes).isAnySuchEdge;
    }
    private ReturnInfo recurHelper(TreeNode node, int totalNodes) {
        ReturnInfo ret = new ReturnInfo(0, false);
        if(node == null) {
            return ret;
        }
        ReturnInfo left = recurHelper(node.left, totalNodes);
        ReturnInfo right = recurHelper(node.right, totalNodes);
        ret.totalCnt = 1 + left.totalCnt + right.totalCnt;
        ret.isAnySuchEdge = ret.totalCnt * 2 == totalNodes || left.isAnySuchEdge || right.isAnySuchEdge;
        return ret;
    }
}