Diameter of a Binary Tree in O(N) in C++

The diameter of a binary tree is the (left_height + right_height + 1) for each node. So in this method we will calculate (left_height + right_height + 1) for each node and update the result . The time complexity here stays O(n).

Let us first define the struct that would represent a tree node that contains the data and its left and right node child. If this is the first node to be created then it’s a root node otherwise a child node.

struct Node {
   int data;
   struct Node *leftChild, *rightChild;
};

Next we create our newNode(int data) function that takes an int value and assign it to the data member of the node. The function returns the pointer to the created struct Node. Also, the left and right child of the newly created node are set to null.

struct Node* newNode(int data){
   struct Node* newNode = new Node;
   newNode->data = data;
   newNode->leftChild = newNode->rightChild = NULL;
   return (newNode);
}

The diameter(Node* root) function takes the root node and checks if the root node is null or not. We then define the ans variable with value INT_MIN. The return value from height(root, ans) is stored to height_of_tree variable. The ans is returned from the function.

int diameter(Node* root){
    if (root == NULL)
        return 0;
    int ans = INT_MIN;
    int height_of_tree = height(root, ans);
    return ans;
}

The height(Node* root, int& ans) function takes the root node and the ans variable by reference. We then perform inorder traversal on the tree to calculate its each subtree length and the maxValue of ans is passed as second parameter on each recursive call. The ans is the max of (ans, 1 + left_height + right_height).

Example

Let us look at the following implementation to find the diameter of the binary tree in O(n) method.

#include <iostream>
using namespace std;
struct Node {
    int data;
    Node* leftChild, *rightChild;
};
struct Node* newNode(int data){
   struct Node* newNode = new Node;
   newNode->data = data;
   newNode->leftChild = newNode->rightChild = NULL;
   return (newNode);
}
int height(Node* root, int& ans){
   if (root == NULL)
      return 0;
   int left_height = height(root->left, ans);
   int right_height = height(root->right, ans);
   ans = max(ans, 1 + left_height + right_height);
   return 1 + max(left_height, right_height);
}
int diameter(Node* root){
   if (root == NULL)
      return 0;
   int ans = INT_MIN;
   int height_of_tree = height(root, ans);
   return ans;
}
int main(){
    struct Node* root = newNode(1);
    root->left = newNode(2);
    root->right = newNode(3);
    root->left->left = newNode(4);
    root->left->right = newNode(5);
    printf("Diameter is %d\n", diameter(root));
    return 0;
}

Output

The above code will produce the following output −

Diameter is 4