Diameter of N-ary Tree

diameter_n_ary_tree.cpp

#include<bits/stdc++.h>
using namespace std;

/*   Let us create below tree
*             A
*         / /  \  \
*       B  F   D  E
*      / \     |  /|\
*     K  J    G  C H I
*      /\            \
*    N   M            L
*/

class TreeNode {
public:
  char val;
  vector<TreeNode*> child;
  TreeNode(char val) : val(val) {}
};

int _diameter(TreeNode* root, int& diameter) {
  if (!root) return 0;
  int leftHeight = 0, rightHeight = 0;

  for (auto child : root->child) {
    int val = _diameter(child, diameter);

    if (leftHeight < val) {
      leftHeight = val;
    }
    else if (rightHeight < val) {
      rightHeight = val;
    }
  }

  diameter = max(diameter, leftHeight + rightHeight + 1);
  return max(leftHeight, rightHeight) + 1;
}

int main() {
  ios_base::sync_with_stdio(false);
  cin.tie(NULL);

  TreeNode* root = new TreeNode('A');
  (root->child).push_back(new TreeNode('B'));
  (root->child).push_back(new TreeNode('F'));
  (root->child).push_back(new TreeNode('D'));
  (root->child).push_back(new TreeNode('E'));
  (root->child[0]->child).push_back(new TreeNode('K'));
  (root->child[0]->child).push_back(new TreeNode('J'));
  (root->child[2]->child).push_back(new TreeNode('G'));
  (root->child[3]->child).push_back(new TreeNode('C'));
  (root->child[3]->child).push_back(new TreeNode('H'));
  (root->child[3]->child).push_back(new TreeNode('I'));
  (root->child[0]->child[0]->child).push_back(new TreeNode('N'));
  (root->child[0]->child[0]->child).push_back(new TreeNode('M'));
  (root->child[3]->child[2]->child).push_back(new TreeNode('L'));
  int diameter = 0;
  _diameter(root, diameter);
  cout << diameter;

}