2023-05-05: Given an undirected, connected tree
There are n nodes labeled 0...n-1 and n-1 edges in the tree.
Given an integer n and an array edges,
edges[i] = [ai, bi] means there is an edge between nodes ai and bi in the tree.
Returns an array answer of length n, where answer[i]:
The sum of the distances between the ith node and all other nodes in the tree.
Input: n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]].
Output: [8,12,6,10,10,10].
Answer 2023-05-05:
Ideas:
Given an undirected, connected tree, it is required to calculate the sum of the distances from each node to all other nodes.
For each node, the sum of its distances to other nodes can be calculated by traversing the tree. For each node, use its child node information to update the sum of its distances to other nodes, and then recursively update its child nodes. Finally, the sum of the distances of all nodes is obtained.
The specific implementation is as follows:
1. Structural diagram
Constructs an undirected graph from the given edgesarray .
2. Traverse the tree and calculate the sum of the distances from each node to other nodes
Recursively traverse the tree starting from the root node, for each node, first initialize the sum of its distances to other nodes to be 0, and then recursively process its child nodes. After processing all child nodes, calculate the sum of the distances from this node to other nodes, and save the size of this node (that is, the number of nodes including itself).
3. Recursively update the sum of the distances from nodes to other nodes
Traverse the tree recursively from the root node, for each node, first calculate the sum of its distances to other nodes, and save it in ansan array . Then recursively process its child nodes, update the sum of their corresponding distances upDistanceinto , and calculate the sum of the distances from each child node to other nodes.
Total time complexity: O(n)
Total space complexity: O(n)
The complete code of go is as follows:
package main
import "fmt"
var N int = 30001
var size [30001]int
var distance [30001]int
func sumOfDistancesInTree(n int, edges [][]int) []int {
graph := make([][]int, n)
for i := range graph {
graph[i] = []int{}
}
for _, edge := range edges {
u := edge[0]
v := edge[1]
graph[u] = append(graph[u], v)
graph[v] = append(graph[v], u)
}
collect(0, -1, graph)
ans := make([]int, n)
setAns(0, -1, 0, graph, ans)
return ans
}
func collect(cur int, father int, graph [][]int) {
size[cur] = 1
distance[cur] = 0
for _, next := range graph[cur] {
if next != father {
collect(next, cur, graph)
distance[cur] += distance[next] + size[next]
size[cur] += size[next]
}
}
}
func setAns(cur int, father int, upDistance int, graph [][]int, ans []int) {
ans[cur] = distance[cur] + upDistance
for _, next := range graph[cur] {
if next != father {
setAns(
next,
cur,
ans[cur]-distance[next]+size[0]-(size[next]<<1),
graph,
ans,
)
}
}
}
func main() {
n := 6
edges := [][]int{{0, 1}, {0, 2}, {2, 3}, {2, 4}, {2, 5}}
result := sumOfDistancesInTree(n, edges)
fmt.Println(result)
}
复制代码The complete code of rust is as follows:
const N: usize = 30001;
static mut SIZE: [i32; N] = [0; N];
static mut DISTANCE: [i32; N] = [0; N];
pub fn sum_of_distances_in_tree(n: i32, edges: Vec<Vec<i32>>) -> Vec<i32> {
let mut graph: Vec<Vec<i32>> = vec![vec![]; n as usize];
for edge in edges {
let u = edge[0] as usize;
let v = edge[1] as usize;
graph[u].push(v as i32);
graph[v].push(u as i32);
}
unsafe {
collect(0, -1, &graph);
let mut ans: Vec<i32> = vec![0; n as usize];
set_ans(0, -1, 0, &graph, &mut ans);
ans
}
}
unsafe fn collect(cur: usize, father: i32, graph: &Vec<Vec<i32>>) {
SIZE[cur] = 1;
DISTANCE[cur] = 0;
for next in &graph[cur] {
let next = *next as usize;
if next != father as usize {
collect(next, cur as i32, graph);
DISTANCE[cur] += DISTANCE[next] + SIZE[next];
SIZE[cur] += SIZE[next];
}
}
}
fn set_ans(cur: usize, father: i32, up_distance: i32, graph: &Vec<Vec<i32>>, ans: &mut Vec<i32>) {
unsafe {
ans[cur] = DISTANCE[cur] + up_distance;
for next in &graph[cur] {
let next = *next as usize;
if next != father as usize {
set_ans(
next,
cur as i32,
ans[cur] - DISTANCE[next] + SIZE[0] - (SIZE[next] << 1),
graph,
ans,
);
}
}
}
}
fn main() {
let n = 6;
let edges = vec![vec![0, 1], vec![0, 2], vec![2, 3], vec![2, 4], vec![2, 5]];
let result = sum_of_distances_in_tree(n, edges);
println!("{:?}", result);
}
复制代码The complete code of c is as follows:
#include <stdio.h>
#include <stdlib.h>
#define N 30001
int size[N];
int distance[N];
void collect(int cur, int father, int** graph, int n);
void setAns(int cur, int father, int upDistance, int** graph, int* ans);
int* sumOfDistancesInTree(int n, int edges[][2]) {
int** graph = malloc(n * sizeof(*graph));
for (int i = 0; i < n; i++) {
graph[i] = malloc((n + 1) * sizeof(**graph));
for (int j = 0; j <= n; j++) {
graph[i][j] = -1;
}
}
for (int i = 0; i < n - 1; i++) {
int u = edges[i][0];
int v = edges[i][1];
if (graph[u][0] == -1) {
graph[u][0] = 0;
}
if (graph[v][0] == -1) {
graph[v][0] = 0;
}
int j = 0;
while (graph[u][++j] != -1);
graph[u][j] = v;
j = 0;
while (graph[v][++j] != -1);
graph[v][j] = u;
}
collect(0, -1, graph, n);
int* ans = malloc(n * sizeof(int));
setAns(0, -1, 0, graph, ans);
for (int i = 0; i < n; i++) {
free(graph[i]);
}
free(graph);
return ans;
}
void collect(int cur, int father, int** graph, int n) {
size[cur] = 1;
distance[cur] = 0;
int j = 1;
while (graph[cur][j] != -1) {
int next = graph[cur][j];
if (next != father) {
collect(next, cur, graph, n);
distance[cur] += distance[next] + size[next];
size[cur] += size[next];
}
j++;
}
}
void setAns(int cur, int father, int upDistance, int** graph, int* ans) {
ans[cur] = distance[cur] + upDistance;
int j = 1;
while (graph[cur][j] != -1) {
int next = graph[cur][j];
if (next != father) {
setAns(
next,
cur,
ans[cur] - distance[next] + size[0] - (size[next] << 1),
graph,
ans
);
}
j++;
}
}
int main() {
int n = 6;
int edges[][2] = { {0, 1}, {0, 2}, {2, 3}, {2, 4}, {2, 5} };
int* result = sumOfDistancesInTree(n, edges);
for (int i = 0; i < n; i++) {
printf("%d ", result[i]);
}
printf("\n");
free(result);
return 0;
}
复制代码The complete c++ code is as follows:
#include <iostream>
#include <vector>
//using namespace std;
const int N = 30001;
static int size[N];
static int distance[N];
void collect(int cur, int father, std::vector<std::vector<int>>& graph);
void setAns(int cur, int father, int upDistance, std::vector<std::vector<int>>& graph, int* ans);
int* sumOfDistancesInTree(int n, std::vector<std::vector<int>>& edges) {
std::vector<std::vector<int>> graph(n);
for (auto edge : edges) {
int u = edge[0];
int v = edge[1];
graph[u].push_back(v);
graph[v].push_back(u);
}
collect(0, -1, graph);
int* ans = new int[n];
setAns(0, -1, 0, graph, ans);
return ans;
}
void collect(int cur, int father, std::vector<std::vector<int>>& graph) {
size[cur] = 1;
distance[cur] = 0;
for (auto next : graph[cur]) {
if (next != father) {
collect(next, cur, graph);
distance[cur] += distance[next] + size[next];
size[cur] += size[next];
}
}
}
void setAns(int cur, int father, int upDistance, std::vector<std::vector<int>>& graph, int* ans) {
int a = N;
ans[cur] = distance[cur] + upDistance;
for (auto next : graph[cur]) {
if (next != father) {
setAns(
next,
cur,
ans[cur] - distance[next] + size[0] - (size[next] << 1),
graph,
ans
);
}
}
}
int main() {
int n = 6;
std::vector<std::vector<int>> edges = { {0, 1}, {0, 2}, {2, 3}, {2, 4}, {2, 5} };
int* result = sumOfDistancesInTree(n, edges);
for (int i = 0; i < n; i++) {
std::cout << result[i] << " ";
}
std::cout << std::endl;
delete[] result;
return 0;
}
复制代码