Height of a general tree in C++

Question

I need to calculate the height of an optional (not AVL, binary) tree. For input I receive an n integer positive number: parent(0), parent(1), ... parent(n-1). Here, parent(i) is a parent of node i. If parent(i) == -1, that i is the root of the tree. It is guaranteed that the sequence has only one root and presets a tree.

Limitation:

\$1 \le n \le 1e5\$

Example input:

5
4 -1 4 1 1

Output:

3

 1
/ \
3  4
   /\
   0 2

My program on C++ works but it looks like that the algorithm is not optimal. The program is verified automatically and one of the test failed with error "Time limit exceeded". Could you advise me, please, on how I can optimize the program?

#include "iostream"

int tree[10000];
int n = 0;

int Height(int parent);

int main ()
{
    std::cin >> n;

    for (int i = 0; i < n; ++i)
    {
        std::cin >> tree[i];
    }

    std::cout << Height(-1) - 1;
    return 0;
}

int Height(int parent)
{
    int height = 1;
    for (int i = 0; i < n; ++i)
    {
        if (tree[i] == parent)
        {
            height = std::max(height, 1 + Height(i));
        }
    }
    return height;
}

Answer 1

Your code is compact, and yet still highly readable. Nice work! Here are some suggestions:

Avoid Global Data

The tree and n variables are global data meaning that any function anywhere in your application can change them. This makes it very difficult to follow the flow of control and figure out where data is modified. It's not a big deal for this little program, but there's no reason you can't put them into main().

Also, while we're talking about your variables, n is not a terribly descriptive name. I would call it something like numNodes or something along those lines.

Don't Traverse the Tree So Many Times

In your Height() function you traverse the tree n^2 times! At the very least, once you've found the parent you shouldn't traverse the rest of the array. But an even better way is to only traverse the entire tree once.

That's what I did in the example below. Note that I couldn't try it out on the website you mentioned because it requires one to sign up for an online course, and as you mentioned, the site was in Russian which I don't understand.

The first thing I did differently was to make a data structure to hold my node data. While the input only contains integers, that doesn't mean that you have to only put integers into your array. I created a type called Node that holds the parent index, a pointer to the actual parent node, and an array of pointers to child nodes. It looks like this:

typedef struct Node {
    int                 parentIndex = { kParentUnset };
    Node*               parent      = { nullptr };
    std::vector<Node*>  children;
} Node;

Which reminds me of my next point - it's a good idea to avoid magic numbers like -1 bare in your code. I made some enums to describe different values for parentIndex:

enum {
    kNoParent       = -1,
    kParentUnset    = -2
};

This makes the code a little easier to read and understand.

You'll notice that I also didn't actually use C arrays, but instead went with std::vector<>. I did this because it allows me to add and remove items without knowing the size ahead of time, and to avoid putting a large block of data on the stack. It's also more standard C++.

So what this code does is read in the number of nodes and the parents of each node. As it reads in a parent, it not only sets the next node's parent index, it also actually sets the pointer to the parent and adds the current node to the parent's list of children. By the time the for loop in main() is complete, we have the entire tree built in memory. The fact is, we don't even need the parent indexes anymore. We could just discard that field, but I've left it to be more clear.

Next, we traverse the tree and calculate the maximum height. This is done using a depth first search. I implemented the search without recursion again to save space on the application stack and avoid a potential stack overflow. We start with the root node and push it onto our local stack variable. Then we enter a loop where we increment the current height, and see if it's greater than the max. If so, we set the max to the current height. Next, we pop the top off the stack. If it has no children, we've hit a leaf node. We decrement the current height and continue on. If it does have children, we push each one onto the stack and loop again.

Here's the code:

int findTreeHeight(Node* rootNode)
{
    int maxHeight   = 0;
    
    std::stack<Node*>   nodeStack;
    nodeStack.push(rootNode);
    
    int currentHeight   = 0;
    while (nodeStack.size() > 0)
    {
        currentHeight++;
        if (currentHeight > maxHeight)
        {
            maxHeight = currentHeight;
        }
        
        Node*   top = nodeStack.top();
        nodeStack.pop();
        
        if (top->children.size() == 0)
        {
            currentHeight--;
        }
        
        for (auto iter = top->children.begin(); iter < top->children.end(); iter++)
        {
            nodeStack.push(*iter);
        }
    }
    
    return maxHeight;
}


int main ()
{
    int numNodes = 0;
    std::cin >> numNodes;

    std::vector<Node> tree(numNodes);
    
    int rootNodeIndex   = -1;
    for (int i = 0; i < numNodes; ++i)
    {
        int parentIndex;
        std::cin >> parentIndex;
        
        tree [ i ].parentIndex = parentIndex;
        
        if (parentIndex == kNoParent)
        {
            rootNodeIndex = i;
        }
        else
        {
            tree [ i ].parent = &tree [ parentIndex ];
            tree [ parentIndex ].children.push_back(&tree [ i ]);
        }
    }
    
    std::cout << findTreeHeight(&tree [ rootNodeIndex ]) << "\n";
    
    return 0;
}

Answer 2

The explanation above is almost excellent. (user1118321) Some cases, this program returns a wrong answer. for example:

if the input will be:

10

9 7 5 5 2 9 9 9 2 -1

               9
           / /  \ \
          / /    \ \
         0  5     6 7
           / \       \
          2   3       1
         / \
        4   8

expected: 4

That solution returns: 5

It can be resolved, replacing this variable "currentHeight" for a height stack which each node has it's respective height.

You can use the second stack with a similar mode the first stack.

I can't share the whole answer, even using this code you will need to implement something else. But I hope this helps. :D

int findTreeHeight(Node* rootNode) {

    int maxHeight   = 0;
    std::stack<Node*>   nodeStack;
    nodeStack.push(rootNode);

    // int currentHeight   = 0; // create the stack here, and add 1; 
    while (nodeStack.size() > 0)
    {
        /* currentHeight++;
        if (currentHeight > maxHeight)
        {
            maxHeight = currentHeight;
        } */

        Node*   top = nodeStack.top();
        nodeStack.pop();

        int currentHeight = stackHeight.top() + 1; // it's a new level
        // after that pop the top of the second stack.  



        if (top->children.size() == 0) {
            // just compare (currentHeight > maxHeight) here because it's a leaf node.       
        }

        for (auto it = top->children.begin(); it < top->children.end(); it++) {
            nodeStack.push(*it);
        }
    }

    return maxHeight;
}

Answer 3

One solution to (user1118321) would be to use a height_stack that keeps track of the height of the nodes in the node_stack, similar to what the answer above suggests.

    for (auto it = top->children.begin(); it < top->children.end(); it++) {
        node_stack.push(*it);
        height_stack.push(current_height+1);
    }

Note: You might have to be careful with how you define current_height for the root (ie. 0 or 1).