# LeetCode 834. Sum of Distances in Tree

https://leetcode.com/problems/sum-of-distances-in-tree/description/

## 題目大意

給 `n` 個樹節點還有表示他們相連情況的 `edge` (沒說是二元樹，就樹而已)
回傳 `ans` 滿足 `ans[i]` 為所有節點到第 `i` 個節點的距離和

## 思考

這題一開始肯定會想要暴力掃一圈去累算答案，但這題 `n` 有點大，用這樣 $O(n^2)$ 肯定不行

我們可以先觀察 example 1

<img src="https://assets.leetcode.com/uploads/2021/07/23/lc-sumdist1.jpg" style="display:block;margin:auto;width:50%;height:auto;" alt="example1"><br>

這題的關鍵就是你要能想像，**我們隨便抓一個點起來，把其他點讓他垂下，就會有樹狀**

然後我們直覺就是會先隨便抓 `0` 當 `root` (因為題目沒有說誰是 `root` ，我們可以不失一般性的抓 `0` 當 `root`)，然後用 DFS 之類的去看其他節點到 `0` 的距離，假設這階段我們取得的答案叫 $\mathrm{dist}(0)$

接下來就是關鍵了，我們現在有 $\mathrm{dist}(0)$ ，以上圖而言，我們接著抓 `2` 當 `root` 時，有沒有辦法用 $\mathrm{dist}(0)$ 的結果去推導 $\mathrm{dist}(2)$ 呢？

我們可以拆兩部分來看：

首先是 `0` 跟 `1` ，他們原本 (以 `0` 當 `root`) 並不是 `2` 的子節點，但現在他們是了，這樣節點每個都相當於比原先距離再多 $1$ (可以想成這些每個都降級，深度增加，因為 `root` 換人當了)

再來是 `2` 的那顆子樹 (`2` - `5`)，他們現在反而是升級了，深度降低，所以每個都少 $1$

如果說 `2` 的子樹有 $m$ 個節點，那麼剩下的屬於我們討論的前者，有 $n-m$ 個節點

於是我們得到：

$$
\mathrm{dist}(2) = \mathrm{dist}(0) + n - 2m
$$

問題現在就簡單了，我們先用 DFS 抓 `0` 為 `root` 掃一輪得到 $\mathrm{dist}(0)$ 跟其他節點的子樹節點數
之後再抓其他節點當 `root` 套用上面的關係式，便求出所求了

class Solution
{
public:
    using Tree = vector<unordered_set<int>>;

public:
    vector<int> sumOfDistancesInTree(int n, vector<vector<int>> &edges)
    {
        this->n = n;
        vector<int> ans(n);
        vector<int> subtreeSize(n, 1);
        Tree tree(n);

        for (const auto &edge : edges)
        {
            const int u = edge[0];
            const int v = edge[1];
            tree[u].insert(v);
            tree[v].insert(u);
        }

        init(tree, 0, -1, subtreeSize, ans);
        reRoot(tree, 0, -1, subtreeSize, ans);
        return ans;
    }

private:
    int n;

    void init(const Tree &tree, int node, int parent, vector<int> &subtreeSize, vector<int> &ans)
    {
        for (const int child : tree[node])
        {
            if (child != parent)
            {
                init(tree, child, node, subtreeSize, ans);
                subtreeSize[node] += subtreeSize[child];
                ans[node] += ans[child] + subtreeSize[child];
            }
        }
    }

    void reRoot(const Tree &tree, int node, int parent, vector<int> &subtreeSize, vector<int> &ans)
    {
        for (const int child : tree[node])
        {
            if (child != parent)
            {
                ans[child] = ans[node] + n - 2 * subtreeSize[child];
                reRoot(tree, child, node, subtreeSize, ans);
            }
        }
    }
};