Number of ways to divide a Binary tree into two halves

Given a binary tree, the task is to count the number of ways to remove a single edge from the tree such that the tree gets divided into two halves with equal sum.
Examples:

Input: 
          1 
        /   \ 
      -1     -1 
               \ 
                1 
Output: 1
Only way to do this will be to remove the edge from the right of the root.
After that we will get 2 sub-trees with sum = 0.
   1 
  /
-1

and 

-1
  \
   1
will be the two sub-trees.

Input:
          1 
        /   \ 
      -1     -1 
               \ 
                -1 
Output: 2

A simple solution will be to remove all the edges of the tree one by one and check if that splits the tree into two halves with the same sum. If it does, we will increase the final answer by 1. This will take O(N²) time in the worst case where “N” is the number of nodes in the tree.
Efficient approach:

Create a variable ‘sum’ and store the sum of all the elements of the Binary tree in it. We can find the sum of all the elements of a Binary tree in O(N) time as discussed in this article.
Now we perform the following steps recursively starting from root node:
- Find the sum of all the elements of its right sub-tree (“R”). If it’s equal to half of the total sum, we increase the count by 1. This is because removing the edge connecting the current node with its right child will divide the tree into two trees with equal sum.
- Find the sum of all the elements of its left sub-tree (“L”). If it’s equal to half of the total sum, we increase the count by 1.

Below is the implementation of the above approach:

C++

// C++ implementation of the approach
#include <bits/stdc++.h>

using namespace std;
 
// Node of a binary tree

struct node {

    int data;

    node* left;

    node* right;

    node(int data)

    {

        this->data = data;

        left = NULL;

        right = NULL;

    }
};
 
// Function to find the sum of
// all the nodes of BST

int findSum(node* curr)
{

    // If current node is

    // null

    if (curr == NULL)

        return 0;
 
    // Else

    return curr->data + findSum(curr->left)

           + findSum(curr->right);
}
 
// Function to recursively check
// if removing any edge divides tree into
// two halves

int checkSum(node* curr, int sum, int& ans)
{

    // Variable to store the

    // sum from left and right

    // child

    int l = 0, r = 0;
 
    // Checking sum from left sub-tree

    // if its not null

    if (curr->left != NULL) {

        l = checkSum(curr->left, sum, ans);

        if (2 * l == sum)

            ans++;

    }
 
    // Checking sum from right sub-tree

    // if its not null

    if (curr->right != NULL) {

        r = checkSum(curr->right, sum, ans);

        if (2 * r == sum)

            ans++;

    }
 
    // Finding the sum of all the elements

    // of current node

    return l + r + curr->data;
}
 
// Function to return the number
// of ways to remove an edge

int cntWays(node* root)
{

    // If root is null

    if (root == NULL)

        return 0;
 
    // To store the final answer

    int ans = 0;
 
    // Sum of all the elements of BST

    int sum = findSum(root);
 
    // If sum is odd then it won't be possible

    // to break it into two halves

    if (sum % 2 == 1)

        return 0;
 
    // Calling the checkSum function

    checkSum(root, sum, ans);
 
    // Returning the final answer

    return ans;
}
 
// Driver code

int main()
{

    node* root = new node(1);

    root->left = new node(-1);

    root->right = new node(-1);

    root->right->right = new node(1);
 
    // Print the count of possible ways

    cout << cntWays(root);
 
    return 0;
}

Java

// Java implementation of the approach

class GFG
{
 
// Node of a binary tree

static class node 
{

    int data;

    node left;

    node right;

    node(int data)

    {

        this.data = data;

        left = null;

        right = null;

    }
};

static int ans;
 
// Function to find the sum of
// all the nodes of BST

static int findSum(node curr)
{

    // If current node is

    // null

    if (curr == null)

        return 0;
 
    // Else

    return curr.data + findSum(curr.left) + 

                       findSum(curr.right);
}
 
// Function to recursively check
// if removing any edge divides tree 
// into two halves

static int checkSum(node curr, int sum)
{

    // Variable to store the

    // sum from left and right

    // child

    int l = 0, r = 0;
 
    // Checking sum from left sub-tree

    // if its not null

    if (curr.left != null) 

    {

        l = checkSum(curr.left, sum);

        if (2 * l == sum)

            ans++;

    }
 
    // Checking sum from right sub-tree

    // if its not null

    if (curr.right != null) 

    {

        r = checkSum(curr.right, sum);

        if (2 * r == sum)

            ans++;

    }
 
    // Finding the sum of all the elements

    // of current node

    return l + r + curr.data;
}
 
// Function to return the number
// of ways to remove an edge

static int cntWays(node root)
{

    // If root is null

    if (root == null)

        return 0;
 
    // To store the final answer

    ans = 0;
 
    // Sum of all the elements of BST

    int sum = findSum(root);
 
    // If sum is odd then it won't be possible

    // to break it into two halves

    if (sum % 2 == 1)

        return 0;
 
    // Calling the checkSum function

    checkSum(root, sum);
 
    // Returning the final answer

    return ans;
}
 
// Driver code

public static void main(String[] args)
{

    node root = new node(1);

    root.left = new node(-1);

    root.right = new node(-1);

    root.right.right = new node(1);
 
    // Print the count of possible ways

    System.out.print(cntWays(root));
}
} 
 
// This code is contributed by PrinciRaj1992

// C# implementation of the approach

using System;
 
class GFG
{
 
// Node of a binary tree

public class node 
{

    public int data;

    public node left;

    public node right;

    public node(int data)

    {

        this.data = data;

        left = null;

        right = null;

    }
};

static int ans;
 
// Function to find the sum of
// all the nodes of BST

static int findSum(node curr)
{

    // If current node is

    // null

    if (curr == null)

        return 0;
 
    // Else

    return curr.data + findSum(curr.left) + 

                       findSum(curr.right);
}
 
// Function to recursively check
// if removing any edge divides tree 
// into two halves

static int checkSum(node curr, int sum)
{

    // Variable to store the

    // sum from left and right

    // child

    int l = 0, r = 0;
 
    // Checking sum from left sub-tree

    // if its not null

    if (curr.left != null) 

    {

        l = checkSum(curr.left, sum);

        if (2 * l == sum)

            ans++;

    }
 
    // Checking sum from right sub-tree

    // if its not null

    if (curr.right != null) 

    {

        r = checkSum(curr.right, sum);

        if (2 * r == sum)

            ans++;

    }
 
    // Finding the sum of all the elements

    // of current node

    return l + r + curr.data;
}
 
// Function to return the number
// of ways to remove an edge

static int cntWays(node root)
{

    // If root is null

    if (root == null)

        return 0;
 
    // To store the final answer

    ans = 0;
 
    // Sum of all the elements of BST

    int sum = findSum(root);
 
    // If sum is odd then it won't be possible

    // to break it into two halves

    if (sum % 2 == 1)

        return 0;
 
    // Calling the checkSum function

    checkSum(root, sum);
 
    // Returning the final answer

    return ans;
}
 
// Driver code

public static void Main(String[] args)
{

    node root = new node(1);

    root.left = new node(-1);

    root.right = new node(-1);

    root.right.right = new node(1);
 
    // Print the count of possible ways

    Console.Write(cntWays(root));
}
} 
 
// This code is contributed by Princi Singh

Javascript

<script>
// javascript implementation of the approach    
// Node of a binary tree
class node {

    constructor(val) {

        this.data = val;

        this.prev = null;

        this.next = null;

    }
}
 
    var ans;
 
    // Function to find the sum of

    // all the nodes of BST

    function findSum( curr) {

        // If current node is

        // null

        if (curr == null)

            return 0;
 
        // Else

        return curr.data + findSum(curr.left) + findSum(curr.right);

    }
 
    // Function to recursively check

    // if removing any edge divides tree

    // into two halves

    function checkSum( curr , sum) {

        // Variable to store the

        // sum from left and right

        // child

        var l = 0, r = 0;
 
        // Checking sum from left sub-tree

        // if its not null

        if (curr.left != null) {

            l = checkSum(curr.left, sum);

            if (2 * l == sum)

                ans++;

        }
 
        // Checking sum from right sub-tree

        // if its not null

        if (curr.right != null) {

            r = checkSum(curr.right, sum);

            if (2 * r == sum)

                ans++;

        }
 
        // Finding the sum of all the elements

        // of current node

        return l + r + curr.data;

    }
 
    // Function to return the number

    // of ways to remove an edge

    function cntWays( root) {

        // If root is null

        if (root == null)

            return 0;
 
        // To store the final answer

        ans = 0;
 
        // Sum of all the elements of BST

        var sum = findSum(root);
 
        // If sum is odd then it won't be possible

        // to break it into two halves

        if (sum % 2 == 1)

            return 0;
 
        // Calling the checkSum function

        checkSum(root, sum);
 
        // Returning the final answer

        return ans;

    }
 
    // Driver code

         root = new node(1);

        root.left = new node(-1);

        root.right = new node(-1);

        root.right.right = new node(1);
 
        // Print the count of possible ways

        document.write(cntWays(root));
 
// This code contributed by Rajput-Ji 
</script>

Python3

# Python3 implementation of the approach
 
# Node of a binary tree

class node:

    def __init__(self,data):

        self.data = data

        self.left = None

        self.right = None
 
# Function to find the sum of
# all the nodes of BST

def findSum(curr):
 
    # If current node is

    # null

    if (curr == None):

        return 0
 
    # Else

    return curr.data + findSum(curr.left)+ findSum(curr.right)
 
# Function to recursively check
# if removing any edge divides tree into
# two halves

def checkSum(curr, s):

    global ans

    # Variable to store the

    # s from left and right

    # child

    l = 0; r = 0
 
    # Checking sum from left sub-tree

    # if its not null

    if (curr.left != None):

        l = checkSum(curr.left, s)

        if (2 * l == s):

            ans+=1
 
    # Checking sum from right sub-tree

    # if its not null

    if (curr.right != None):

        r = checkSum(curr.right, s)

        if (2 * r == s):

            ans+=1
 
    # Finding the sum of all the elements

    # of current node

    return l + r + curr.data
 
# Function to return the number
# of ways to remove an edge

def cntWays(root):

    # If root is null

    if (root == None):

        return 0
 
    # To store the final answer

    global ans

    ans = 0
 
    # s of all the elements of BST

    s = findSum(root)
 
    # If s is odd then it won't be possible

    # to break it into two halves

    if (s % 2):

        return 0
 
    # Calling the checkSum function

    checkSum(root, s)
 
    # Returning the final answer

    return ans
 
# Driver code

if __name__ == '__main__':
 
    root = node(1)

    root.left = node(-1)

    root.right = node(-1)

    root.right.right = node(1)
 
    # Print the count of possible ways

    print(cntWays(root))

Output:

Time complexity of this approach will be O(N) and space complexity will O(H) where “N” equals number of node in Binary tree and “H” equals height of the Binary Tree.

Article Tags :

Data Structures

DSA

Tree

Binary Tree