diff options
-rw-r--r-- | 834/go.mod | 3 | ||||
-rw-r--r-- | 834/main.go | 51 |
2 files changed, 54 insertions, 0 deletions
diff --git a/834/go.mod b/834/go.mod new file mode 100644 index 0000000..2a95dd6 --- /dev/null +++ b/834/go.mod @@ -0,0 +1,3 @@ +module 834 + +go 1.19 diff --git a/834/main.go b/834/main.go new file mode 100644 index 0000000..c2392f0 --- /dev/null +++ b/834/main.go @@ -0,0 +1,51 @@ +package main + +import "fmt" + +func dfs(node int, parent int, adjList map[int][]int, count []int, result []int) { + for _, child := range adjList[node] { + if child == parent { + continue + } + + dfs(child, node, adjList, count, result) + count[node] += count[child] + result[node] += result[child] + count[child] + } +} + +func dfs2(node int, parent int, adjList map[int][]int, count []int, result []int) { + for _, child := range adjList[node] { + if child == parent { + continue + } + + result[child] = result[node] - count[child] + len(count) - count[child] + dfs2(child, node, adjList, count, result) + } +} + +func sumOfDistancesInTree(n int, edges [][]int) []int { + adjList := make(map[int][]int) + + for _, edge := range edges { + adjList[edge[0]] = append(adjList[edge[0]], edge[1]) + adjList[edge[1]] = append(adjList[edge[1]], edge[0]) + } + + count := make([]int, n) + for i := 0; i < len(count); i++ { + count[i] = 1 + } + + result := make([]int, n) + + dfs(0, -1, adjList, count, result) + dfs2(0, -1, adjList, count, result) + + return result +} + +func main() { + fmt.Println("vim-go") +} |