Leetcode.com problem 834: Sum of Distances in Tree

James S. Plank

Sat Dec 24 01:22:33 EST 2022

In case Leetcode's servers are down

You are to write a method in the class Solution with the following prototype: 

  vector<int> sumOfDistancesInTree(int n, vector <vector <int> >&edges);

Here's an example tree (which is example 1): 

    0
   / \
  /   \
 1    2
     /|\
    / | \
   3  4  5

The answer is { 8, 12, 6, 10, 10, 10 }. To explain a little, consider node 0. Its distance to nodes 1 and 2 is one, and its distance to nodes 3, 4 and 5 is two. So the sum of its distances is 1+1+2+2+2 = 8.

Constraints: The tree has up to 30,000 nodes.

The skeleton program

If you give it no command line arguments, then it will read the edges on standard input. If you give it a positive number on the command line, then it will generate a random problem. If the number is a multiple of 5, the problem will be small. Otherwise, the problem will be of a random size less than or equal to 30,000. If you give it a negative number, then it will print the edges before calling the solution. For example: 
UNIX> a.out 5                // When the seed is a multiple of 5, the tree is small.
15 21 13 19 23 17 21 13 
UNIX> a.out -5               // When you give it a negative number, it prints the edges.
5 2         
7 2
0 7
4 5
3 2
1 0
6 0

15 21 13 19 23 17 21 13 
UNIX>

Hints

It's not a bad idea to work through the example above before moving on. The first thing to realize is that any node can be the root of the tree. It doesn't affect the answer. To illustrate, suppose we make node 5 the root of the tree: 

   5
   |
   2
  /|\
 / | \
0  3  4 
|
1

You'll note that the answer is the same. So let's do a few things to start:

I don't know how to help you think about constructing an answer for problems like these. Your first thought should be: With 30,000 nodes, my solution has to be O(n) or O(n log n) at the worst. DFS, BFS and tree traversals are all O(n) so let's start there. What quantities can we store in each node that we can set with one of those algorithms?

Task 1 -- Assigning Parents

The first, most obvious one is assigning parents of each node. You can do that with a tree traversal. Make that your first task. To help you test, I have two files. ex1.txt has the tree above. ex4.txt has the tree below: 

        0
       /|\
      / | \
     1  3  2 
    /|    / \
   / |   /   \
  4  8   14   15
 /|\    /|    /|\
/ | \  / |   / | \
5 6 7  9 10 /  |  \
           11  12 13

I added a Parent vector to the Solution class, and then wrote a recursive method to set the parents. When traversing a node's adjacency list, I simply skipped the parent and didn't make a recursive call on the parent (I do this in all of my tree traversals). 

void set_parents(int node);

I tested it and put the outputs into parents1.txt and parents4.txt. Test yourself with these files. 

parents1.txt

0: -1
1: 0
2: 0
3: 2
4: 2
5: 2

parents4.txt

0: -1
1: 0
2: 0
3: 0
4: 1
5: 4
6: 4
7: 4
8: 1
9: 14
10: 14
11: 15
12: 15
13: 15
14: 2
15: 2

Task 2 -- Setting Sizes

A second simple task is to create a vector Size and setting Size[i] to the number of nodes in the tree rooted by node i. This is a post-order traversal. Here's output for the two trees: 

sizes1.txt

0: 6
1: 1
2: 4
3: 1
4: 1
5: 1

sizes4.txt

0: 16
1: 6
2: 8
3: 1
4: 4
5: 1
6: 1
7: 1
8: 1
9: 1
10: 1
11: 1
12: 1
13: 1
14: 3
15: 4

Task 3 -- sum of distances to the nodes in your tree

Next, I tried to think of something I could do with a postorder traversal. The task I came up with is to sum all of the distances of a node to nodes in the tree that it roots. For example, for node 4 in the ex4.txt, this sum is three. For node 1, the sum is 8 (Distance 1 to nodes 4 and 8. Distance 2 to nodes 5, 6 and 7.)

How do we do this with a postorder traversal? Work through the example of node 1 above. The sums for nodes 5, 6, 7 and 8 are zero. Node 4 has a sum of 3. When we calculate the sum for node one, we can do it as follows:

Go ahead and write this code. Here are the answers for the two trees. I called the array Below, since it is the sum of distances to nodes below a given node. If you are having problems with this, walk though how you calculate the values for nodes 14, 15, 2 and 0 above. 

below1.txt

0: 8
1: 0
2: 3
3: 0
4: 0
5: 0

below4.txt

0: 35
1: 8
2: 12
3: 0
4: 3
5: 0
6: 0
7: 0
8: 0
9: 0
10: 0
11: 0
12: 0
13: 0
14: 2
15: 3

Task 4 -- calculating the answer

First, note that Below[0] is in fact the answer for node 0. That's because all of the nodes are in node 0's tree.

Now, let's consider node 1 in ex4.txt. I'm going to redraw the tree in two colors. 

        0
       /|\
      / | \
     1  3  2 
    /|    / \
   / |   /   \
  4  8   14   15
 /|\    /|    /|\
/ | \  / |   / | \
5 6 7  9 10 /  |  \
           11  12 13

We can calculate the answer for node 1 by counting the sum of distances to the blue nodes, and then adding the sum of distances to the red nodes. The sum of distances to the blue nodes is simple -- it's Below[1].

To calculate the red nodes, let's count the sum of distances from node 0 to the red nodes. That's going to be the following equation:

SumRed0 = Answer[0] - (Below[1] + Size[1])

The sum of distances from node 1 to the red nodes (which we'll call SumRed1) will be SumRed0 plus the number of red nodes. What is that number? It is: (Number of nodes in the tree) - Size[1]). So: SumRed1 = SumRed0 + #Nodes - Size[1]. Put it all together:

Similarly:

You do this with a preorder traversal. Here are the answers: 

answers1.txt

0: 8
1: 12
2: 6
3: 10
4: 10
5: 10

answers4.txt

0: 35
1: 39
2: 35
3: 49
4: 47
5: 61
6: 61
7: 61
8: 53
9: 59
10: 59
11: 57
12: 57
13: 57
14: 45
15: 43

I hope that gives you all of the pieces to the puzzle, so that you can solve the problem!