Count the number of Nodes in a Binary Tree in Constant Space

Given a binary tree having N nodes, count the number of nodes using constant O(1) space. This can be done by simple traversals like- Preorder, InOrder, PostOrder, and LevelOrder but these traversals require an extra space which is equal to the height of the tree.

Examples:

Input:

Output: 5
Explanation: In the above binary tree, there are total 5 nodes: 1, 2, 3, 4 and 5

Input:

Output: 7
Explanation: In the above binary tree, there are total 7 nodes: 1, 2, 3, 4, 5, 6 and 7

Approach: To solve the problem follow the below idea:

We can solve this problem using the Morris Traversal—a Threaded Binary Tree traversal technique that enables us traverse all the nodes in any binary tree in constant space. Morris traversal does not rely on any additional data structure. The process operates by altering the arrangement of the tree as we navigate through it and subsequently restoring it back to the original state.

Steps to count nodes in a tree using Morris Traversal:

• Start by initializing two variables, cnt as 0 to keep track of the number of nodes counted and current as root to traverse the array.
• Repeat the following steps until we have traversed the tree:
  • If the current node does not have a left child increase the count. Move to its right child.
  • If the current node has a left child find its inorder predecessor (the rightmost node, in its subtree).
  • If the predecessors right child is null set it to be the node and move to its child.
  • If the predecessors right child is already set as the node (indicating that we have visited its subtree) reset it to null increment the count and move to its right child.
• Finally return the count variable which will contain the number of nodes in this tree.

Below is the implementation for the above approach:

C++

// C++ Implementation

#include <iostream>

class TreeNode {
public:
    int val;
    TreeNode* left;
    TreeNode* right;

    TreeNode(int value)
    {
        val = value;
        left = nullptr;
        right = nullptr;
    }
};

int countNodes(TreeNode* root)
{
    int count = 0;
    TreeNode* current = root;

    while (current != nullptr) {
        if (current->left == nullptr) {
            // If the left child is null, it's a leaf node,
            // so increment the count and move to the right
            // child.
            count++;
            current = current->right;
        }
        else {
            // If there is a left child, find the
            // predecessor (the rightmost node in the left
            // subtree).
            TreeNode* predecessor = current->left;
            while (predecessor->right != nullptr
                   && predecessor->right != current) {
                predecessor = predecessor->right;
            }

            if (predecessor->right == nullptr) {
                // If the predecessor's right child is null,
                // establish a link and move to the left
                // child.
                predecessor->right = current;
                current = current->left;
            }
            else {
                // If the predecessor's right child is not
                // null, remove the link, increment the
                // count, and move to the right child.
                predecessor->right = nullptr;
                count++;
                current = current->right;
            }
        }
    }

    return count;
}

int main()
{
    TreeNode* root = new TreeNode(1);
    root->left = new TreeNode(2);
    root->right = new TreeNode(3);
    root->left->left = new TreeNode(4);
    root->left->right = new TreeNode(5);

    int nodeCount = countNodes(root);

    std::cout << "Number of nodes in the binary tree: "
              << nodeCount << std::endl;

    return 0;
}

// This code is contributed by Sakshi

Java

class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;

    TreeNode(int value) {
        val = value;
    }
}

public class CountNodesInBinaryTree {
    public static int countNodes(TreeNode root) {
        int count = 0;
        TreeNode current = root;

        while (current != null) {
            if (current.left == null) {
                // If the left child is null, it's a leaf node, so increment the count and move to the right child.
                count++;
                current = current.right;
            } else {
                // If there is a left child, find the predecessor (the rightmost node in the left subtree).
                TreeNode predecessor = current.left;
                while (predecessor.right != null && predecessor.right != current) {
                    predecessor = predecessor.right;
                }

                if (predecessor.right == null) {
                    // If the predecessor's right child is null, establish a link and move to the left child.
                    predecessor.right = current;
                    current = current.left;
                } else {
                    // If the predecessor's right child is not null, remove the link, increment the count, and move to the right child.
                    predecessor.right = null;
                    count++;
                    current = current.right;
                }
            }
        }

        return count;
    }

    public static void main(String[] args) {
        TreeNode root = new TreeNode(1);
        root.left = new TreeNode(2);
        root.right = new TreeNode(3);
        root.left.left = new TreeNode(4);
        root.left.right = new TreeNode(5);

        int nodeCount = countNodes(root);

        System.out.println("Number of nodes in the binary tree: " + nodeCount);
    }
}

Python

# Python code for the above approach

class TreeNode:
    def __init__(self, value):
        self.val = value
        self.left = None
        self.right = None

def countNodes(root):
    cnt = 0
    current = root

    while current is not None:
        if current.left is None:
            cnt += 1
            current = current.right
        else:
            pre = current.left
            while pre.right is not None and pre.right != current:
                pre = pre.right

            # If right child doesn't point to the
            # current node, then update the
            # right child pointer with the current
            # node (add the thread) and update
            # the current with it's left child
            if pre.right is None:
                pre.right = current
                current = current.left
            else:
            # If the right child of inorder
            # predecessor already points to the
            # current node, (remove the thread)
            # increment the count by one and update
            # the current with it's right child
                pre.right = None
                cnt += 1
                current = current.right

    return cnt

root = TreeNode(1)
root.left = TreeNode(2)
root.right = TreeNode(3)
root.left.left = TreeNode(4)
root.left.right = TreeNode(5)

# Count the nodes in the binary tree using the countNodes function
nodeCount = countNodes(root)

# Print the result
print('Number of nodes in the binary tree:', nodeCount)

C#

using System;

public class TreeNode
{
    public int val;         // Value of the node
    public TreeNode left;   // Reference to the left child node
    public TreeNode right;  // Reference to the right child node

    // Constructor to initialize the TreeNode with a value
    public TreeNode(int value)
    {
        val = value;
        left = null;  // Initializing left child as null
        right = null; // Initializing right child as null
    }
}

public class BinaryTree
{
    // Function to count the number of nodes in the binary tree
    public int CountNodes(TreeNode root)
    {
        int count = 0;        // Initialize the count of nodes to zero
        TreeNode current = root; // Start traversal from the root node

        // Loop until the current node being checked is not null
        while (current != null)
        {
            if (current.left == null)
            {
                // If the left child is null, it's a leaf node,
                // so increment the count and move to the right child.
                count++;             // Increment the count of nodes
                current = current.right; // Move to the right child
            }
            else
            {
                // If there is a left child, find the predecessor
                // (the rightmost node in the left subtree).
                TreeNode predecessor = current.left;

                // Traverse right to find the predecessor
                while (predecessor.right != null &&
                       predecessor.right != current)
                {
                    predecessor = predecessor.right;
                }

                if (predecessor.right == null)
                {
                    // If the predecessor's right child is null,
                    // establish a link and move to the left child.
                    predecessor.right = current; // Link predecessor to current node
                    current = current.left;      // Move to the left child
                }
                else
                {
                    // If the predecessor's right child is not null,
                    // remove the link, increment the count, and move
                    // to the right child.
                    predecessor.right = null; // Remove the link to restore the tree structure
                    count++;                  // Increment the count of nodes
                    current = current.right;  // Move to the right child
                }
            }
        }

        return count; // Return the total count of nodes in the binary tree
    }

    public static void Main(string[] args)
    {
        // Create a sample binary tree
        TreeNode root = new TreeNode(1);
        root.left = new TreeNode(2);
        root.right = new TreeNode(3);
        root.left.left = new TreeNode(4);
        root.left.right = new TreeNode(5);

        BinaryTree binaryTree = new BinaryTree(); // Create an instance of BinaryTree
        int nodeCount = binaryTree.CountNodes(root); // Get the count of nodes

        // Display the number of nodes in the binary tree
        Console.WriteLine("Number of nodes in the binary tree: " + nodeCount);
    }
}

Javascript

// JavaScript Implementation
class TreeNode {
  constructor(value) {
    this.val = value;
    this.left = null;
    this.right = null;
  }
}

function countNodes(root) {
  let count = 0;
  let current = root;

  while (current !== null) {
    if (current.left === null) {
      // If the left child is null, it's a leaf node,
      // so increment the count and move to the right child.
      count++;
      current = current.right;
    }
    else {
      // If there is a left child, find the predecessor
      // (the rightmost node in the left subtree).
      let predecessor = current.left;
      while (predecessor.right !== null && predecessor.right !== current) {
        predecessor = predecessor.right;
      }

      if (predecessor.right === null) {
        // If the predecessor's right child is null,
        // establish a link and move to the left child.
        predecessor.right = current;
        current = current.left;
      }
      else {
        // If the predecessor's right child is not null,
        // remove the link, increment the count, and move
        // to the right child.
        predecessor.right = null;
        count++;
        current = current.right;
      }
    }
  }

  return count;
}

// Driver Code
const root = new TreeNode(1);
root.left = new TreeNode(2);
root.right = new TreeNode(3);
root.left.left = new TreeNode(4);
root.left.right = new TreeNode(5);

const nodeCount = countNodes(root);

console.log("Number of nodes in the binary tree: " + nodeCount);

Output

('Number of nodes in the binary tree:', 5)

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