How to prune tree data structure at given depth in C - c

Question

I am trying to write this function:
struct treeNode *pruneTree(struct treeNode *root, int depth);
Which given a tree like:
1 level 0
/ \
2 3 level 1
/ \ \
4 5 6 level 2
/ \
7 8 level 3
If depth = 1 then create a tree with depth = 1 and cut everything after so the result should be:
1
/ \
2 3 // tree with depth = 1
I know how to write a function that prunes the leaves and I am trying to adapt it to prune at any level:
int isLeaf (struct treeNode * treeNode) {
return (treeNode->left == NULL) && (treeNode->right == NULL);
}
void removeLeaves(struct treeNode * root) {
if (root->left != NULL) {
if (isLeaf (root->left)) {
free(root->left);
}
else {
removeLeaves(root->left);
}
}
if (root->right != NULL) {
if (isLeaf (root->right)) {
free(root->right);
}
else {
removeLeaves(root->right);
}
}
}
What is a good strategy to do this? My approach is to replace the isLeaf function with a isAfterDepth function and using a helper function that calculates the depth, but this doesn't seem efficient. What is a more elegant way to do it?

Answer 1

Copying the tree
If you intend to make a copy of the tree that is pruned at a certain level, you can simply use recursion and at each recursive call decrease the depth parameter with one, if the depth results in 0, you simply do no longer recursively copy the children.
struct treeNode *pruneTree(struct treeNode *root, int depth) { //version where the tree is copied
if(root == NULL || depth < 0) {
return NULL;
} else {
treeNode *copyRoot = (treeNode*) malloc(sizeof(treeNode));
copyRoot->value = root->value;
copyRoot->left = pruneTree(root->left,depth-1);
copyRoot->right = pruneTree(root->right,depth-1);
return copyRoot;
}
}
The code works as follows: if the given root pointer is NULL or th depth is less than zero, NULL is returned, because either we call this with the child of a leaf or the depth constraint has been reached.
If that is not the case, we make a copy of the current node: we allocate a new treeNode object, copy the value of the original node (assuming this is called value), and perform a recursive call to copy the left and right children.
Altering the tree
You can also alter the current tree. In that case, you better first define a function to remove a subtree and all its descendants:
void prune(struct treeNode * root) {
if(root != NULL) {
if (root->left != NULL) {
prune(root->left);
}
if (root->right != NULL) {
prune(root->right);
}
free(root);
}
}
Now we simply need to define a method that will prune only at a certain level:
struct treeNode *pruneTree(struct treeNode *root, int depth) { //version where the tree is altered
if(root != NULL) {
if(depth <= 0) {
prune(root->left);
root->left = NULL;
prune(root->right);
root->right = NULL;
} else {
pruneTree(root->left,depth-1);
pruneTree(root->right,depth-1);
}
}
return root;
}