Open In App

Count of nodes with average of left subtree at least K in a given Binary Tree

Improve
Improve
Like Article
Like
Save
Share
Report

Given a binary tree and a number K, the task is to count the number of nodes having the average of the values in their left subtree greater than or equal to K.

Examples:

Input :  K=5
Tree:    
               2
           /     \
        5        4
      / \       /  \
   5   6    6   2
    \       /
    5    4
Output: 3
Explanation:  
               2         ——– level 0
           /     \                        
        5        4       ——– level 1
      / \       /  \                    
   5   6      6   2     ——– level 2
    \         /                          
    5      4            ——– level 3

Left left subtree average for root node 2 = (5+5+5) / 3 = 5
Left left subtree average for node 4 at level 1 = (6+4) / 2 = 5
Left left subtree average for node 5 at level 1 = (5+5) / 2 = 5
So, these 3 nodes satisfy the given conditions.

Input: K = 4,  
Tree:           1
                /
              2
            /
          3
        /
      4
Output: 1

 

Method 1 :

Approach: Follow the steps below to solve this problem:

  1. Create a global variable ans to store the answer and initialise it with 0.
  2. Create a function countHelper which will accept a tree node and integer K as the parameter, and will return a pair of the sum of nodes and number of nodes in the subtree of that node.
  3. Now, pass the root node and K in the initial call of this function.
  4. In each call of this recursive function:
    1. Check if the current node is a leaf node or not. If it is a leaf node, just return {0, 0} as both the sum and number of nodes below this node is 0.
    2. Now, call for left and right subtrees.
    3. Check if for the current node, (the sum of left and right subtree/ number of nodes in left and right subtree)>=K, if it is increment ans by 1.
  5. After the recursive function ends, print ans.

Below is the implementation of the above approach.

C++




// C++ code for the above approach
#include <bits/stdc++.h>
using namespace std;
 
// Structure of a tree node
class Node {
public:
    int data;
    Node *left, *right;
    Node(int d)
    {
        data = d;
        left = right = NULL;
    }
};
 
// Global variable to store the node count
int ans = 0;
 
// Function to count the nodes in a tree
// with average of all left nodes
// greater than or equal to K .
pair<int, int> countHelper(Node* root,
                           int K)
{
 
    // For leaf Node
    if (!root->left and !root->right) {
        return { root->data, 1 };
    }
 
    pair<int, int> left = { 0, 0 },
                   right = { 0, 0 };
 
    // For left subtree
    if (root->left) {
        left = countHelper(root->left, K);
    }
 
    // For right subtree
    if (root->right) {
        right = countHelper(root->right, K);
    }
    if (left.second != 0
        and left.first / left.second >= K) {
        ans += 1;
    }
 
    return { left.first + right.first + root->data,
             left.second + right.second + 1 };
}
 
// Function to call the initial
// instance of function countHelper
void countNodes(Node* root, int K)
{
    countHelper(root, K);
    cout << ans;
}
 
// Driver Code
int main()
{
    // Given Tree
    Node* root = new Node(2);
    root->left = new Node(5);
    root->right = new Node(4);
    root->left->left = new Node(5);
    root->left->right = new Node(6);
    root->right->left = new Node(6);
    root->right->right = new Node(2);
    root->left->left->right = new Node(5);
    root->right->left->left = new Node(4);
 
    int K = 5;
 
    countNodes(root, K);
    return 0;
}


Java




// Java code implementation of above approach
import java.util.*;
 
public class Solution {
 
    static class Pair {
        int first, second;
        public Pair(int first, int second) {
            this.first = first;
            this.second = second;
        }
    }
 
    // Structure of a tree node
    static class Node {
        int data;
        Node left, right;
        Node(int d) {
            data = d;
            left = right = null;
        }
    };
 
    // Global variable to store the node count
    static int ans = 0;
 
    // Function to count the nodes in a tree
    // with average of all left nodes
    // greater than or equal to K.
    static Pair countHelper(Node root, int K) {
        // For leaf Node
        if (root.left == null && root.right == null) {
            return new Pair(root.data, 1);
        }
 
        Pair left = new Pair(0, 0);
        Pair right = new Pair(0, 0);
 
        // For left subtree
        if (root.left != null) {
            left = countHelper(root.left, K);
        }
 
        // For right subtree
        if (root.right != null) {
            right = countHelper(root.right, K);
        }
 
        if (left.second != 0 && left.first / left.second >= K) {
            ans += 1;
        }
 
        return new Pair(left.first + right.first + root.data,
                        left.second + right.second + 1);
    }
 
    // Function to call the initial
    // instance of function countHelper
    static void countNodes(Node root, int K) {
        countHelper(root, K);
        System.out.print(ans);
    }
 
    // Driver Code
    public static void main(String[] args) {
        // Given Tree
        Node root = new Node(2);
        root.left = new Node(5);
        root.right = new Node(4);
        root.left.left = new Node(5);
        root.left.right = new Node(6);
        root.right.left = new Node(6);
        root.right.right = new Node(2);
        root.left.left.right = new Node(5);
        root.right.left.left = new Node(4);
 
        int K = 5;
 
        countNodes(root, K);
    }
}


Python3




# Python code for the above approach
class pair:
    def __init__(self, first, second):
        self.first = first;
        self.second = second;
 
# Structure of a tree node
class Node:
    def __init__(self, d):
        self.data = d;
        self.left = self.right = None;
 
# Global variable to store the node count
ans = 0;
 
# Function to count the nodes in a tree
# with average of all left nodes
# greater than or equal to K .
def countHelper(root, K):
 
    # For leaf Node
    if (root.left == None and root.right == None):
        return pair(root.data, 1);
 
    left = pair(0, 0)
    right = pair(0, 0)
 
    # For left subtree
    if (root.left != None):
        left = countHelper(root.left, K);
 
    # For right subtree
    if (root.right != None):
        right = countHelper(root.right, K);
     
    if (left.second != 0 and left.first / left.second >= K):
        global ans;
        ans += 1;
 
    return pair(left.first + right.first + root.data, left.second + right.second + 1);
 
# Function to call the initial
# instance of function countHelper
def countNodes(root, K):
    countHelper(root, K);
    print(ans);
 
# Driver Code
# Given Tree
root = Node(2);
root.left = Node(5);
root.right = Node(4);
root.left.left = Node(5);
root.left.right = Node(6);
root.right.left = Node(6);
root.right.right = Node(2);
root.left.left.right = Node(5);
root.right.left.left = Node(4);
 
K = 5;
 
countNodes(root, K);
 
# This code is contributed by Saurabh Jaiswal.


C#




// C# code for the above approach
using System;
using System.Collections.Generic;
 
public class GFG{
 
  class pair
  {
    public int first, second;
    public pair(int first, int second) 
    {
      this.first = first;
      this.second = second;
    }   
  }
 
  // Structure of a tree node
  class Node {
    public int data;
    public Node left, right;
    public Node(int d)
    {
      data = d;
      left = right = null;
    }
  };
 
  // Global variable to store the node count
  static int ans = 0;
 
  // Function to count the nodes in a tree
  // with average of all left nodes
  // greater than or equal to K .
  static pair countHelper(Node root,
                          int K)
  {
 
    // For leaf Node
    if (root.left==null && root.right==null) {
      return new pair( root.data, 1 );
    }
 
    pair left = new pair( 0, 0 ),
    right = new pair( 0, 0 );
 
    // For left subtree
    if (root.left!=null) {
      left = countHelper(root.left, K);
    }
 
    // For right subtree
    if (root.right!=null) {
      right = countHelper(root.right, K);
    }
    if (left.second != 0
        && left.first / left.second >= K) {
      ans += 1;
    }
 
    return new pair( left.first + right.first + root.data,
                    left.second + right.second + 1 );
  }
 
  // Function to call the initial
  // instance of function countHelper
  static void countNodes(Node root, int K)
  {
    countHelper(root, K);
    Console.Write(ans);
  }
 
  // Driver Code
  public static void Main(String[] args)
  {
 
    // Given Tree
    Node root = new Node(2);
    root.left = new Node(5);
    root.right = new Node(4);
    root.left.left = new Node(5);
    root.left.right = new Node(6);
    root.right.left = new Node(6);
    root.right.right = new Node(2);
    root.left.left.right = new Node(5);
    root.right.left.left = new Node(4);
 
    int K = 5;
 
    countNodes(root, K);
  }
}
 
// This code is contributed by shikhasingrajput


Javascript




<script>
// Javascript code for the above approach
 
 
class pair {
    constructor(first, second) {
        this.first = first;
        this.second = second;
    }
}
 
// Structure of a tree node
class Node {
 
    constructor(d) {
        this.data = d;
        this.left = this.right = null;
    }
};
 
// Global variable to store the node count
let ans = 0;
 
// Function to count the nodes in a tree
// with average of all left nodes
// greater than or equal to K .
function countHelper(root, K) {
 
    // For leaf Node
    if (root.left == null && root.right == null) {
        return new pair(root.data, 1);
    }
 
    let left = new pair(0, 0),
        right = new pair(0, 0);
 
    // For left subtree
    if (root.left != null) {
        left = countHelper(root.left, K);
    }
 
    // For right subtree
    if (root.right != null) {
        right = countHelper(root.right, K);
    }
    if (left.second != 0
        && left.first / left.second >= K) {
        ans += 1;
    }
 
    return new pair(left.first + right.first + root.data, left.second + right.second + 1);
}
 
// Function to call the initial
// instance of function countHelper
function countNodes(root, K) {
    countHelper(root, K);
    document.write(ans);
}
 
// Driver Code
// Given Tree
let root = new Node(2);
root.left = new Node(5);
root.right = new Node(4);
root.left.left = new Node(5);
root.left.right = new Node(6);
root.right.left = new Node(6);
root.right.right = new Node(2);
root.left.left.right = new Node(5);
root.right.left.left = new Node(4);
 
let K = 5;
 
countNodes(root, K);
 
// This code is contributed by Saurabh Jaiswal
</script>


Output

3






Time Complexity: O(N) where N is the number of nodes in the tree
Auxiliary Space: O(N)

Method 2: Using Depth-First Search (DFS)

In the below code approaches, the tree can be traversed using a depth-first search while recording the total value and the number of nodes in the left subtree for each node. Next, we may determine whether the average of the left subtree is larger than or equal to K for each node, and if so, increase the count.

Approach: Follow the steps below to solve this problem:

  1. Define the Node structure for the binary tree.
  2. Define a function ‘dfs‘ which will take the root node of the binary tree, a reference to an integer variable ‘ans‘, a reference to an integer variable ‘sum‘, a reference to an integer variable ‘count’, and an integer variable ‘K’ as input.
  3. If the root is NULL, return.
  4. Call ‘dfs’ function recursively with the left child of the current node as the root, and update the ‘left_sum’, ‘left_count’ variables with the sum of values and count of nodes in the left subtree respectively.
  5. Call ‘dfs’ function recursively with the right child of the current node as the root, and update the ‘right_sum’, ‘right_count’ variables with the sum of values and count of nodes in the right subtree respectively.
  6. Calculate the ‘sum’ variable as the sum of ‘left_sum’, ‘right_sum’ and the current node’s data.
  7. Calculate the ‘count’ variable as the sum of ‘left_count’, ‘right_count‘ and 1.
  8. If ‘left_count’ is not 0 and the average of left subtree values (‘left_sum’ / ‘left_count’) is greater than or equal to ‘K’, then increment the ‘ans’ variable.
  9. Define the ‘countNodes‘ function which will take the root node of the binary tree and an integer variable ‘K’ as input.
  10. Initialize integer variables ‘ans’, ‘sum’ and ‘count’ to 0.
  11. Call ‘dfs’ function with the root node of the binary tree, ‘ans’, ‘sum’, ‘count’, and ‘K’ as inputs.
  12. Print the value of ‘ans’ which represents the number of nodes whose left subtree average is greater than or equal to ‘K’.
  13. In the ‘main’ function, create the binary tree as mentioned in the problem statement.
  14. Set the value of ‘K’ as 5.
  15. Call ‘countNodes’ function with the root node of the binary tree and ‘K’ as inputs.
  16. Return 0.

Below is the implementation of the above approach.

C++




// C++ code implementation for the above approach
#include <bits/stdc++.h>
using namespace std;
 
// Structure of a tree node
class Node {
public:
    int data;
    Node *left, *right;
    Node(int d)
    {
        data = d;
        left = right = NULL;
    }
};
 
// Function to perform DFS traversal
void dfs(Node* root, int& ans, int& sum, int& count, int K) {
    if(!root) return;
     
    int left_sum = 0, left_count = 0;
    dfs(root->left, ans, left_sum, left_count, K);
     
    int right_sum = 0, right_count = 0;
    dfs(root->right, ans, right_sum, right_count, K);
     
    sum = left_sum + right_sum + root->data;
    count = left_count + right_count + 1;
     
    if(left_count != 0 && left_sum / left_count >= K) {
        ans++;
    }
}
 
// Function to count the nodes in a tree
// with average of all left nodes
// greater than or equal to K.
void countNodes(Node* root, int K)
{
    int ans = 0, sum = 0, count = 0;
    dfs(root, ans, sum, count, K);
    cout << ans << endl;
}
 
// drive Code
int main()
{
    // Given Tree
    Node* root = new Node(2);
    root->left = new Node(5);
    root->right = new Node(4);
    root->left->left = new Node(5);
    root->left->right = new Node(6);
    root->right->left = new Node(6);
    root->right->right = new Node(2);
    root->left->left->right = new Node(5);
    root->right->left->left = new Node(4);
 
    int K = 5;
 
    countNodes(root, K);
    return 0;
}


Java




// Java Code
 
import java.io.*;
 
// Structure of a tree node
class Node {
    int data;
    Node left, right;
 
    Node(int d) {
        data = d;
        left = right = null;
    }
}
 
public class Program {
    // Function to perform DFS traversal
    public static void dfs(Node root, int[] ans, int[] sum, int[] count, int K) {
        if (root == null) return;
 
        int[] leftSum = {0};
        int[] leftCount = {0};
        dfs(root.left, ans, leftSum, leftCount, K);
 
        int[] rightSum = {0};
        int[] rightCount = {0};
        dfs(root.right, ans, rightSum, rightCount, K);
 
        sum[0] = leftSum[0] + rightSum[0] + root.data;
        count[0] = leftCount[0] + rightCount[0] + 1;
 
        if (leftCount[0] != 0 && leftSum[0] / leftCount[0] >= K) {
            ans[0]++;
        }
    }
 
    // Function to count the nodes in a tree
    // with the average of all left nodes
    // greater than or equal to K.
    public static void countNodes(Node root, int K) {
        int[] ans = {0};
        int[] sum = {0};
        int[] count = {0};
        dfs(root, ans, sum, count, K);
        System.out.println(ans[0]);
    }
 
    // Drive Code
    public static void main(String[] args) {
        // Given Tree
        Node root = new Node(2);
        root.left = new Node(5);
        root.right = new Node(4);
        root.left.left = new Node(5);
        root.left.right = new Node(6);
        root.right.left = new Node(6);
        root.right.right = new Node(2);
        root.left.left.right = new Node(5);
        root.right.left.left = new Node(4);
 
        int K = 5;
 
        countNodes(root, K);
    }
}


Python3




# Define a Node class to represent a binary tree node.
class Node:
    def __init__(self, d):
        self.data = d
        self.left = None
        self.right = None
 
# Define a depth-first search function to traverse the binary tree.
# This function calculates the number of nodes whose average value of the subtree is greater than or equal to K.
def dfs(root, ans, sum, count, K):
    if root is None:
        return
 
    # Initialize variables to store the sum and count for the left and right subtrees.
    leftSum = [0]
    leftCount = [0]
 
    # Recursively call the dfs function on the left subtree.
    dfs(root.left, ans, leftSum, leftCount, K)
 
    # Initialize variables to store the sum and count for the right subtree.
    rightSum = [0]
    rightCount = [0]
 
    # Recursively call the dfs function on the right subtree.
    dfs(root.right, ans, rightSum, rightCount, K)
 
    # Calculate the sum and count for the current node and its subtrees.
    sum[0] = leftSum[0] + rightSum[0] + root.data
    count[0] = leftCount[0] + rightCount[0] + 1
 
    # Check if the average of the left subtree is greater than or equal to K and increment the answer.
    if leftCount[0] != 0 and leftSum[0] / leftCount[0] >= K:
        ans[0] += 1
 
# Define a function to count nodes in the binary tree whose average value of the subtree is greater than or equal to K.
def countNodes(root, K):
    ans = [0# Initialize a variable to store the answer.
    sum = [0# Initialize a variable to store the sum.
    count = [0# Initialize a variable to store the count.
    dfs(root, ans, sum, count, K)  # Call the dfs function to calculate the answer.
    print(ans[0])  # Print the final answer.
 
# Drive Code
# Given Tree
root = Node(2)
root.left = Node(5)
root.right = Node(4)
root.left.left = Node(5)
root.left.right = Node(6)
root.right.left = Node(6)
root.right.right = Node(2)
root.left.left.right = Node(5)
root.right.left.left = Node(4)
 
K = 5
 
countNodes(root, K)  # Call the countNodes function to find and print the answer.


C#




using System;
 
// Structure of a tree node
public class Node
{
    public int data;
    public Node left, right;
 
    public Node(int d)
    {
        data = d;
        left = right = null;
    }
}
 
public class Program
{
    // Function to perform DFS traversal
    public static void dfs(Node root, ref int ans, ref int sum, ref int count, int K)
    {
        if (root == null) return;
 
        int left_sum = 0, left_count = 0;
        dfs(root.left, ref ans, ref left_sum, ref left_count, K);
 
        int right_sum = 0, right_count = 0;
        dfs(root.right, ref ans, ref right_sum, ref right_count, K);
 
        sum = left_sum + right_sum + root.data;
        count = left_count + right_count + 1;
 
        if (left_count != 0 && left_sum / left_count >= K)
        {
            ans++;
        }
    }
 
    // Function to count the nodes in a tree
    // with average of all left nodes
    // greater than or equal to K.
    public static void countNodes(Node root, int K)
    {
        int ans = 0, sum = 0, count = 0;
        dfs(root, ref ans, ref sum, ref count, K);
        Console.WriteLine(ans);
    }
 
    // Drive Code
    public static void Main()
    {
        // Given Tree
        Node root = new Node(2);
        root.left = new Node(5);
        root.right = new Node(4);
        root.left.left = new Node(5);
        root.left.right = new Node(6);
        root.right.left = new Node(6);
        root.right.right = new Node(2);
        root.left.left.right = new Node(5);
        root.right.left.left = new Node(4);
 
        int K = 5;
 
        countNodes(root, K);
    }
}


Javascript




// Structure of a tree node
class Node {
    constructor(d) {
        this.data = d;
        this.left = null;
        this.right = null;
    }
}
 
// Function to perform DFS traversal
function dfs(root, ans, sum, count, K) {
    if (root === null) return;
 
    const leftSum = [0];
    const leftCount = [0];
    dfs(root.left, ans, leftSum, leftCount, K);
 
    const rightSum = [0];
    const rightCount = [0];
    dfs(root.right, ans, rightSum, rightCount, K);
 
    sum[0] = leftSum[0] + rightSum[0] + root.data;
    count[0] = leftCount[0] + rightCount[0] + 1;
 
    if (leftCount[0] !== 0 && leftSum[0] / leftCount[0] >= K) {
        ans[0]++;
    }
}
 
// Function to count the nodes in a tree
// with the average of all left nodes
// greater than or equal to K.
function countNodes(root, K) {
    const ans = [0];
    const sum = [0];
    const count = [0];
    dfs(root, ans, sum, count, K);
    console.log(ans[0]);
}
 
// Drive Code
// Given Tree
const root = new Node(2);
root.left = new Node(5);
root.right = new Node(4);
root.left.left = new Node(5);
root.left.right = new Node(6);
root.right.left = new Node(6);
root.right.right = new Node(2);
root.left.left.right = new Node(5);
root.right.left.left = new Node(4);
 
const K = 5;
 
countNodes(root, K);


Output

3







Time Complexity : O(N)
Auxiliary Space: O(N)



Last Updated : 17 Nov, 2023
Like Article
Save Article
Previous
Next
Share your thoughts in the comments
Similar Reads