leetcode

0 class Solution {

1 public:

2 vector<int> sumOfDistancesInTree(int n, vector<vector<int>> &edges) {

3 vector<vector<int>> adj(n);

4 vector<int> res(n), count(n, 1);

5

6 for (const auto &edge : edges) {

7 adj[edge[0]].push_back(edge[1]);

8 adj[edge[1]].push_back(edge[0]);

9 }

10

11 using record_t = tuple<int, int>;

12 stack<record_t> st;

13

14 for (st.emplace(-1, 0); !st.empty();) {

15 if (get<1>(st.top()) != -1) {

16 const auto [parent, root] = st.top();

17 st.emplace(-1, -1);

18 for (const auto next : adj[root]) {

19 if (next == parent) continue;

20 st.emplace(root, next);

21 }

22 continue;

23 }

24

25 st.pop();

26 const auto [parent, root] = st.top();

27 st.pop();

28

29 for (const auto next : adj[root]) {

30 if (next == parent) continue;

31 count[root] += count[next];

32 res[root] += res[next] + count[next];

33 }

34 }

35

36 for (st.emplace(-1, 0); !st.empty();) {

37 const auto [parent, root] = st.top();

38 st.pop();

39

40 for (const auto next : adj[root]) {

41 if (next == parent) continue;

42 res[next] = res[root] - count[next] + (n - count[next]);

43 st.emplace(root, next);

44 }

45 }

46

47 return res;

48 }

49 };