Leet Code Daily 2024-04-28 - Sum of Distances in Tree

Description

Solution

Let's take node 0 as the root and form the graph as a tree.

By DFS, we can get:

$C(i) = \text{set of children of node } i$
$P(i) = \text{parent of node } i$

If we define:

$F(i) = \text{number of nodes in sub-tree } i$
$G(i) = \text{sum of distances from node } i \text{ to all decents of node } i$
$H(i) = \text{sum of distances from node } i \text{ to all non-decents of node } i$

Obviously,

F(i) = \sum_{c \in C(i)}{F(c) + 1}

G(i) = \sum_{c \in C(i)}{F(c) + G(c)}

And for $H(i)$ , considering both parent node and sibling nodes, we can find:

H(i) = \begin{cases} 0 & i = 0 \\ H(P(i)) + N - F(P(i)) + 1 + \sum_{j \in Child(P(i)) \land j \neq i }{G(j) + 2 \times F(j)} & i \neq 0 \\ \end{cases}

Also we know:

\begin{split} \sum_{j \in Child(P(i)) \land j \neq i }{G(j) + 2 \times F(j)} & = G(P(i)) - G(i) - F(i) + F(P(i)) - F(i) - 1 \\ & = G(P(i)) + F(P(i)) - G(i) - 2 \times F(i) - 1 \end{split}

Thus:

H(i) = \begin{cases} 0 & i = 0 \\ H(P(i)) + N + G(P(i)) - G(i) - 2 \times F(i) & i \neq 0 \\ \end{cases}

And the answer of node i would be $G(i) + H(i)$ .

Code

c++

constexpr int kNotComputed = -1;
constexpr int kRoot = 0;

class Solution {
public:
    vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
        // Build adjacency.
        vector<vector<int>> adjacency(n);
        for (const auto edge : edges) {
            adjacency[edge[0]].push_back(edge[1]);
            adjacency[edge[1]].push_back(edge[0]);
        }

        // DFS to compute P(i) & C(i).
        vector<int> P(n, kNotComputed);
        vector<vector<int>> C(n);
        function<void(int)> dfs;
        dfs = [&](int i) {
            for (const auto j : adjacency[i]) {
                if (P[j] == kNotComputed && j != kRoot) {
                    P[j] = i;
                    C[i].push_back(j);
                    dfs(j);
                }
            }
        };
        dfs(kRoot);

        // Compute F(i), G(i) & H(i).
        vector<int> F(n, kNotComputed), G(n, kNotComputed), H(n, kNotComputed);
        function<int(int)> getF, getG, getH;
        getF = [&](int i) -> int {
            if (F[i] == kNotComputed) {
                F[i] = 1;
                for (const auto c : C[i]) {
                    F[i] += getF(c);
                }
            }
            return F[i];
        };
        getG = [&](int i) -> int {
            if (G[i] == kNotComputed) {
                G[i] = 0;
                for (const auto c : C[i]) {
                    G[i] += getF(c) + getG(c);
                }
            }
            return G[i];
        };
        getH = [&](int i) -> int {
            if (H[i] == kNotComputed) {
                H[i] = (i == kRoot) ? 0
                                    : (H[i] = getH(P[i]) + n + getG(P[i]) -
                                              getG(i) - 2 * getF(i));
            }
            return H[i];
        };

        // Build answer.
        vector<int> answer;
        for (int i = 0; i < n; i++) {
            answer.push_back(getG(i) + getH(i));
        }
        return answer;
    }
};

Leet Code Daily 2024-04-28 - Sum of Distances in Tree

Solution

Code

COMMAND INTERFACE

:help aguang.xyz