Count nodes with specific number of children in a binary tree? - c++

Question

I have this challenging exercise I got from a book about c++, and i'm not sure how to tackle this problem.
I must define a function called treeNodeCount() which returns the number of nodes in a binary tree (easy enough), and I also have to define an overloaded function that takes an int(0,1, or 2) which represents the number of children, and the function should return the nodes that have that specific number of children.
treeNodeCount should both use a function called nodeCount(elemType root) to do the recursion necessary to count the nodes(so basically all the work).
challenge number one says that I can add a second parameter to nodeCount which takes the number of children for the nodes that we want to count.
Challenge number two says that we cannot use a second parameter (this is the tough part)
I was able to do challenge one and here is what I came up with:
template <class elemType>
int binaryTreeType<elemType>::nodeCount(nodeType<elemType> *p, int a ) const
{
if (p == NULL){
return 0;
}
else if (a == 0 && p->lLink == NULL && p->rLink == NULL){
return 1 + nodeCount(p->lLink, a) + nodeCount(p->rLink, a);
}
else if (a == 1 && (p->lLink != NULL ^ p->rLink != NULL)){
return 1 + nodeCount(p->lLink, a) + nodeCount(p->rLink, a);
}
else if (a == 2 && p->lLink != NULL && p->rLink != NULL){
return 1 + nodeCount(p->lLink, a) + nodeCount(p->rLink, a);
}
else if (a == -1){
return nodeCount(p->lLink, a) + nodeCount(p->rLink, a) + 1;
}
return nodeCount(p->lLink, a) + nodeCount(p->rLink, a);
}
template <class elemType>
int binaryTreeType<elemType>::treeNodeCount(int a) const{
return nodeCount(root, a);
}
This seems to work fine but i am convinced that there has to be a better way.
I was not able to do challenge 2 though, and i have no idea what to do (is it even possible)

Answer 1

You can scrunch down your logic and make it a bit more straightforward by implementing a function to return the number of children given a node.
template <class elemType>
int nodeSize(nodeType<elemType>* node) const
{
int count = 0;
if (node->lLink)
++count;
if (node->rLink)
++count;
return count;
}
template <class elemType>
int binaryTreeType<elemType>::nodeCount(nodeType<elemType>* node, int count) const
{
if (node)
{
if (nodeSize(node) == count || count == -1)
return nodeCount(node->lLink, count) + nodeCount(node->rLink, count) + 1;
return nodeCount(node->lLink, count) + nodeCount(node->rLink, count);
}
return 0;
}
For the second challenge, you need a stack to avoid recursion.
template <class elemType>
int binaryTreeType<elemType>::treeNodeCount(int count) const
{
stack<nodeType<elemType>*> node_stack;
node_stack.push(root);
int num_matches = 0;
while (!stack.empty())
{
nodeType<elemType>* node = node_stack.top();
node_stack.pop();
if (node)
{
if (nodeSize(node) == count || count == -1)
++num_matches;
node_stack.push(node->lLink);
node_stack.push(node->rLink);
}
}
return num_matches;
}
Edit: fixed a goof in the above recursive version. Thanks to David Rodriguez for pointing it out.