Counting unival subtrees

Question

I have written some code to solve the following interview question. Please advise how it can be improved. Thanks in advance.

A unival tree (which stands for "universal value") is a tree where all nodes under it have the same value. Given the root to a binary tree, count the number of unival subtrees. For example, the following tree has 5 unival subtrees:

    0
   / \
  1   0
     / \
    1   0
   / \
  1   1

Implementation:

#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>

typedef struct stTree
{
    struct stTree * left;
    struct stTree * right;
    int value;
}

stTree;

stTree* createNode(int value)
{
    stTree *node = malloc(sizeof *node);
    node->left = NULL;
    node->right = NULL;
    node->value = value;

    return node;
}

bool isTreeUniv(stTree *node)
{
    bool flag = true;

    if (!node)
        return false;

    if (node->right && node->right->value != node->value)
    {
        flag = false;
    }

    if (node->left && node->left->value != node->value)
    {
        flag = false;
    }

    return flag;
}

stTree* insertRight(stTree *currNode, int value)
{
    stTree *node = malloc(sizeof *node);

    currNode->right = node;
    node->left = NULL;
    node->right = NULL;
    node->value = value;
    return node;
}

stTree* insertLeft(stTree *currNode, int value)
{
    stTree *node = malloc(sizeof *node);

    currNode->left = node;
    node->left = NULL;
    node->right = NULL;
    node->value = value;
    return node;
}

unsigned int uTreeCount = 0;
void countUnivSubT(stTree *Node)
{
    if (isTreeUniv(Node))
        uTreeCount++;

    if (Node->left)
        countUnivSubT(Node->left);

    if (Node->right)
        countUnivSubT(Node->right);

}

int main(void)
{
    //build a tree
    stTree *rootNode = createNode(0);
    insertLeft(rootNode, 1);
    insertRight(rootNode, 0);

    insertLeft(rootNode->right, 1);
    insertRight(rootNode->right, 0);

    insertLeft(rootNode->right->left, 1);
    insertRight(rootNode->right->left, 1);

    countUnivSubT(rootNode);
    printf("total universal subree: %u\n", uTreeCount);

}

Answer 1

Code-formatting

1. Putting the tag-type two lines down instead of on the same line as the closing brace is curious.

Managing trees

2. I wonder why you don't use createNode() in insertRight() and insertLeft().

3. If you change to building the trees from the leaves down to the root instead the other way around, you only need a single createNode() accepting a value and two (possibly NULL) descendants.

4. Assuming that resource-aquisition always succeeds is quite brave.

5. Consider adding a way to free a tree, for best effect using constant space, even though using it just before tearing down the whole process is unconscionably wasteful.

stTree* createNode(int value, stTree* left, stTree* right) {
    stTree* r = malloc(sizeof *r);
    if (!r) abort();
    r->value = value;
    r->left = left;
    r->right = right;
    return r;
}

static stTree* findBottomLeft(stTree* node) {
    while (node->left)
        node = node->left;
    return node;
}
void freeTree(stTree* node) {
    if (!node) return;
    stTree* bottomLeft = findBottomLeft(node);
    while (node) {
        if (node->right) {
            bottomLeft->left = node->right;
            bottomLeft = findBottomLeft(bottomLeft);
        }
        stTree* old = node;
        node = node->left;
        free(old);
    }
}

The main part

6. If you don't need to modify something, don't require the right. Use const.

7. isTreeUniv() is just broken. It only checks the direct descendents, while it should recurse into them.

8. Consequently, countUnivSubT() is also wrong. Still, fixing isTreeUniv() would result in a \$O(n^2)\$ algorithm, when it should be \$O(n)\$. The idea is to get all the info you need at once.

9. Avoid globals. Using uTreeCount makes the code non-reentrant, and breaks locality of reasoning.

static bool countUnivSubTimpl(const stTree* node, const stTree* parent, size_t* count) {
    if (!node) return true;
    bool r = countUnivSubTimpl(node->left, node, count)
        & countUnivSubTimpl(node->right, node, count);
    *count += r;
    return r & node->value == parent->value;
}
size_t countUnivSubT(const stTree* node) {
    size_t r = 0;
    countUnivSubTimpl(node, node, &r);
    return r;
}

int main() {
    stTree* root =
        createNode(0,
            createNode(1, NULL, NULL),
            createNode(0,
                createNode(1,
                    createNode(1, NULL, NULL),
                    createNode(1, NULL, NULL),
                ),
                createNode(0, NULL, NULL)));

    size_t uTreeCount = countUnivSubT(root);
    printf("total universal subree: %u\n", uTreeCount);
}

Answer 2

The one red flag is uTreeCount. This should not be a global, and in fact it is easy to rephrase your countUnivSubT to be fully re-entrant: have it return an integer, and do addition within the body, something like

unsigned countUnivSubT(stTree *Node)
{
    unsigned int uTreeCount = isTreeUniv(Node);

    if (Node->left)
        uTreeCount += countUnivSubT(Node->left);

    if (Node->right)
        uTreeCout += countUnivSubT(Node->right);

    return uTreeCount;
}

That said, you have an inner null check, so this can actually reduce to

unsigned countUnivSubT(stTree *Node)
{
    if (!Node) return 0;
 
    return isTreeUniv(Node)
        + countUnivSubT(Node->left)
        + countUnivSubT(Node->right);
}

Answer 3

const

Consider const with functions that do not alter the tree:

//  bool isTreeUniv(stTree *node)
bool isTreeUniv(const stTree *node)

//void countUnivSubT(stTree *Node)
void countUnivSubT(const stTree *Node)

This improves clarity of what code does and allows for select optimizations.

Loop opportunity vs recursion

Rather than a global and two recursive calls, perhaps loop on one side and recurse on the other:

unsigned countUnivSubT(const stTree *Node) {
  unsigned count = 0;
  while (Node) {
    count += isTreeUniv(Node);
    if (Node->left) {
      count += countUnivSubT(Node->left);
    } 
    Node = Node->right;
  }
  return count;
}

Answer 4

The algorithm is inefficient.

For each node, we examine all nodes in its subtree to determine whether they are all equal. We should look to visit each node just once, and extract as much as we need in that single visit. So, as we go, report back up whether the current node is a unival tree, as well as the count of unival trees at or below it. We don't need to visit the children again, just use the retrieved information. Like this:

static size_t countUnivSubT_impl(const stTree *node, bool *isUnival, int *value)
{
    if (!node) {
        return 0;
    }
    *value = node->value;

    /* initial values chosen to work if one/both children are null */
    int lval = node->value, rval = node->value;
    bool lunival = true, runival = true;

    size_t count_left = countUnivSubT_impl(node->left, &lunival, &lval);
    size_t count_right = countUnivSubT_impl(node->right, &runival, &rval);
    return count_left + count_right
        +  (*isUnival = /* N.B. assignment */
            lunival && lval == node->value &&
            runival && rval == node->value);
}

size_t countUnivSubT(const stTree *node)
{
    bool isUnival;
    int value;
    return countUnivSubT_impl(node, &isUnival, &value);
}

And use it in main():

printf("There are %zu universal subtrees\n",
       countUnivSubT(rootNode));

(I corrected the spelling there, too).