1
/*
2

3
    Finding sum of distance from ith node to all other nodes takes O(n), so it
4
   will take O(n^2) if done naively to find distance for all.
5

6
            x---------------------y
7
          /  \                    \
8
         o    o                    o
9
        /                         / \
10
       o                         o   o
11

12
      Above is a graph where x and y are subtrees with them as the root. They
13
   both are neighbors and are connected by an edge.
14

15
      From x's POV
16
             x
17
          /  |   \
18
         o    o   y
19
        /          \
20
       o            o
21
                  /  \
22
                 o    o
23

24
    From y's POV
25
                 y
26
               /    \
27
              x      o
28
            /  |    /  \
29
           o   o   o    o
30
          /
31
         o
32

33
    As evident froma above, the tree structure looks different when we change
34
   the root of tree. So even if we find the sum of distance for node 0 as root,
35
   the subanswers are actually not correct for other nodes when they are the
36
   root of tree. The subanswers are only correct when node 0 is root, since that
37
   is how it was computed.
38

39
    We use something called re-rooting, this allows us to find a relation so
40
   that we can change the root of the tree and use the answer from its neighbor
41
   to compute answer for it. So if we compute the answer for node 0 as root, we
42
   can find the answers for its children, then those children can be used to
43
   compute for their neighbors and so on.
44

45
    For the 1st diagram:
46
    distance[x] = sum(x) + sum(y) + count(y)
47

48
    distance[x] is the overall sum of distance with root as x
49
    sum(i) is the sum of distance of all the descendents for a subtree rooted at
50
   i count(i) is the no. of descendent nodes in the subtree rooted at i
51

52
    Now why count(y) ?
53
    Consider a node z in subtree y, dist(x, z) = dist(y, z) + 1
54
    So sum(y) already accounts for dist(y, z) and we just need to add 1
55
    So if there are n_y nodes in the subtree, we need to add +1 that many times.
56

57
    distance[x] = sum(x) + sum(y) + count(y) -----------------1
58
    distance[y] = sum(y) + sum(x) + count(x) -----------------2
59

60
    From (1) - (2)
61
    distance[x] - distance[y] = count(y) - count(x) ----------3
62
    distance[x] = distance[y] - count(x) + count(y)
63

64
    Above relation can be used to find the answer for a neighbor, when the
65
   answer for the other neighbor is known.
66

67
    In our case, we can compute the answer for node 0 as root in one traversal
68
   of post order. Then in another traversal with again root as 0, compute the
69
   answer for its children and from them to their children nodes and so on.
70

71
    distance[child] = distance[parent] - count(child) + (N - count(child))
72
    Since we also know the no. of nodes in subtree with child as root, the
73
   remaining nodes = (N - count(child))
74

75

76
    Re-rooting Ref:
77
   https://leetcode.com/problems/sum-of-distances-in-tree/solution/
78

79
    TC: O(N)
80
    SC: O(N)
81
*/
82
class Solution {
83
 private:
84
  int n = 0;
85
  // subtree_distance[i] = Sum of distance to other nodes when the ith-node is
86
  // root
87
  vector<int> subtree_distance;
88
  // subtree_count[i] = no. of nodes in the subtree with ith-node as root
89
  vector<int> subtree_count;
90

91
 public:
92
  void postorder(int root, int parent, vector<vector<int>> &graph) {
93
    // Perform DFS for the child nodes
94
    for (auto child : graph[root]) {
95
      // Avoid iterating to the parent, it will create a loop otherwise
96
      if (child != parent) {
97
        postorder(child, root, graph);
98
        // Update the subtree count and sum of distance
99
        subtree_count[root] += subtree_count[child];
100
        // distance[X] = distance[X] + distance[Y] + n_Y
101
        subtree_distance[root] += subtree_distance[child] + subtree_count[child];
102
      }
103
    }
104
  }
105

106
  void preorder(int root, int parent, vector<vector<int>> &graph) {
107
    for (auto child : graph[root]) {
108
      if (child != parent) {
109
        // distance[child] = distance[parent] - count[child] + count(parent)
110
        subtree_distance[child] =
111
            subtree_distance[root] - subtree_count[child] + (n - subtree_count[child]);
112
        preorder(child, root, graph);
113
      }
114
    }
115
  }
116

117
  vector<int> reRootingSol(int n, vector<vector<int>> &edges) {
118
    // create an undirected graph
119
    vector<vector<int>> graph(n);
120
    for (auto edge : edges) {
121
      int src = edge[0], dst = edge[1];
122
      graph[src].emplace_back(dst);
123
      graph[dst].emplace_back(src);
124
    }
125

126
    this->n = n;
127
    this->subtree_count.resize(n, 1);
128
    this->subtree_distance.resize(n, 0);
129

130
    // This computes the subtree sum and subtree node count when 0 is the root
131
    // of graph Imagine looking at the graph from node 0's POV
132
    postorder(0, -1, graph);
133
    // Since we have computed the sum distance from node's POV, we can use that
134
    // information to find the sum of distance from each node's POV i.e imagine
135
    // looking at the graph from POV of each node
136
    preorder(0, -1, graph);
137
    return this->subtree_distance;
138
  }
139

140
  vector<int> sumOfDistancesInTree(int n, vector<vector<int>> &edges) {
141
    return reRootingSol(n, edges);
142
  }
143
};