Сложность: hard
Есть неориентированное связное дерево с n узлами, пронумерованными от 0 до n - 1, и n - 1 ребрами.
Вам даны целое число n и массив edges, где edges[i] = [ai, bi] указывает, что существует ребро между узлами ai и bi в дереве.
Верните массив answer длиной n, где answer[i] — это сумма расстояний между i-м узлом в дереве и всеми другими узлами.
Пример:
Input: n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
Output: [8,12,6,10,10,10]
Explanation: The tree is shown above.
We can see that dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
equals 1 + 1 + 2 + 2 + 2 = 8.
Hence, answer[0] = 8, and so on.
👨💻 Алгоритм:
1⃣Постройте дерево и выполните обход в глубину (DFS) для расчета количества узлов в поддереве и суммы расстояний до всех узлов поддерева.
2⃣На основе полученных данных рассчитайте сумму расстояний для корня дерева.
3⃣Выполните второй обход в глубину (DFS) для корректировки суммы расстояний для каждого узла на основе суммы расстояний его родительского узла.
😎 Решение:
class Solution {
var ans: [Int] = []
var count: [Int] = []
var graph: [[Int]] = []
var N: Int = 0
func sumOfDistancesInTree(_ N: Int, _ edges: [[Int]]) -> [Int] {
self.N = N
graph = Array(repeating: [], count: N)
ans = Array(repeating: 0, count: N)
count = Array(repeating: 1, count: N)
for edge in edges {
graph[edge[0]].append(edge[1])
graph[edge[1]].append(edge[0])
}
dfs(0, -1)
dfs2(0, -1)
return ans
}
func dfs(_ node: Int, _ parent: Int) {
for child in graph[node] {
if child != parent {
dfs(child, node)
count[node] += count[child]
ans[node] += ans[child] + count[child]
}
}
}
func dfs2(_ node: Int, _ parent: Int) {
for child in graph[node] {
if child != parent {
ans[child] = ans[node] - count[child] + N - count[child]
dfs2(child, node)
}
}
}
}Ставь 👍 и забирай 📚 Базу знаний