给定一个无向、连通的树

树中有 n 个标记为 0…n-1 的节点以及 n-1 条边 。

给定整数 n 和数组 edges ，

edges[i] = [ai, bi]表示树中的节点 ai 和 bi 之间有一条边。

返回长度为 n 的数组 answer ，其中 answer[i] :

树中第 i 个节点与所有其他节点之间的距离之和。

输入: n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]。

输出: [8,12,6,10,10,10]。

思路：

给定一棵无向、连通的树，要求计算每个节点到其他所有节点的距离之和。

可以通过遍历树，对于每个节点分别计算它到其他节点的距离之和。对于每个节点，利用它的子节点信息来更新它到其他节点的距离之和，然后递归地更新它的子节点。最终得到所有节点的距离之和。

具体实现如下：

1.构造图

通过给定的 edges 数组构造无向图。

2.遍历树，计算每个节点到其他节点的距离之和

从根节点开始递归遍历树，对于每个节点，首先初始化它到其他节点的距离之和为 0，然后递归地处理它的子节点。处理完所有子节点之后，计算该节点到其他节点的距离之和，并将该节点的大小（即包括自身在内的节点数）保存下来。

3.递归更新节点到其他节点的距离之和

从根节点开始递归遍历树，对于每个节点，首先计算它到其他节点的距离之和，并将其保存在 ans 数组中。然后递归地处理它的子节点，将它们对应的距离之和更新到 upDistance 中，并计算每个子节点到其他节点的距离之和。

总时间复杂度：O(n)

总空间复杂度：O(n)

go完整代码如下：

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)
}

rust完整代码如下：

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);
}

c完整代码如下：

#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;
}

c++完整代码如下：

#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;
}