LeetCode Entry
834. Sum of Distances in Tree
fun sumOfDistancesInTree(n: Int, edges: Array
834. Sum of Distances in Tree hard
https://t.me/leetcode_daily_unstoppable/60
fun sumOfDistancesInTree(n: Int, edges: Array<IntArray>): IntArray {
val graph = mutableMapOf<Int, MutableList<Int>>()
edges.forEach { (from, to) ->
graph.getOrPut(from, { mutableListOf() }) += to
graph.getOrPut(to, { mutableListOf() }) += from
}
val counts = IntArray(n) { 1 }
val sums = IntArray(n) { 0 }
fun distSum(pos: Int, visited: Int) {
graph[pos]?.forEach {
if (it != visited) {
distSum(it, pos)
counts[pos] += counts[it]
sums[pos] += counts[it] + sums[it]
}
}
}
fun dfs(pos: Int, visited: Int) {
graph[pos]?.forEach {
if (it != visited) {
sums[it] = sums[pos] - counts[it] + (n - counts[it])
dfs(it, pos)
}
}
}
distSum(0, -1)
dfs(0, -1)
return sums
}
We can do the job for item #0, then we need to invent a formula to reuse some data when we change the node.
How to mathematically prove formula for a new sum:
Store count of children in a counts array, and sum of the distances to children in a dist array. In a first DFS traverse from a node 0 and fill the arrays. In a second DFS only modify dist based on previous computed dist value, using formula: sum[curr] = sum[prev] - count[curr] + (N - count[curr])
Space: O(N), Time: O(N)