« Count Number of Nodes in a tree

TreeNode.cpp

Copy
1#include <vector>
2using namespace std;
3
4template <typename T>
5class TreeNode
6{
7public:
8    T data;
9    vector<TreeNode<T> *> children;
10
11    TreeNode(T data)
12    {
13        this->data = data;
14    }
15
16    // current node won't be deleted before destructor is completed. So here we can do the cleanup
17    ~TreeNode()
18    {
19        for (int i = 0; i < children.size(); i++)
20        {
21            delete children[i];
22        }
23    }
24};

TreeUse.cpp

Copy
1#include <iostream>
2#include <iomanip>
3#include <algorithm>
4#include <string>
5#include <cstring>
6#include <vector>
7#include <cmath>
8#include <map>
9#include <climits>
10// climits for INT_MIN
11#include <unordered_map>
12using namespace std;
13#include "TreeNode.h"
14#include <queue>
15
16TreeNode<int> *takeInputLevelWise()
17{
18    int rootData;
19    cout << "Enter root data: " << endl;
20    cin >> rootData;
21    TreeNode<int> *root = new TreeNode<int>(rootData);
22
23    queue<TreeNode<int> *> pendingNodes;
24
25    pendingNodes.push(root);
26    while (!pendingNodes.empty())
27    {
28        TreeNode<int> *front = pendingNodes.front();
29        pendingNodes.pop();
30        cout << "Enter number of children of " << front->data << endl;
31        int numChild;
32        cin >> numChild;
33
34        for (int i = 0; i < numChild; i++)
35        {
36            int childData;
37            cout << "Enter " << i << "th child of " << front->data << endl;
38            cin >> childData;
39            TreeNode<int> *child = new TreeNode<int>(childData);
40            front->children.push_back(child);
41            pendingNodes.push(child);
42        }
43    }
44    return root;
45}
46
47int numNodes(TreeNode<int> *node)
48{
49    if (node == NULL) // not needed but to handle edge cases
50    {
51        return 0;
52    }
53    int ans = 1;
54    for (int i = 0; i < node->children.size(); i++)
55    {
56        ans += numNodes(node->children[i]);
57    }
58    return ans;
59}
60
61int main()
62{
63    TreeNode<int> *newRoot = takeInputLevelWise();
64    printTree(newRoot);
65
66    int ans = numNodes(newRoot);
67    cout << "NUM NODES:" << ans << endl;
68
69    return 0;
70}