LeetCode 834. 树中距离之和

Categories:

题目描述
思路分析
1. 解法一：树形 DP + 换根（推荐）
相似题目

题目描述

✅ 834. 树中距离之和

思路分析

解法一：树形 DP + 换根（推荐）

核心思路：

第一次 DFS 计算每个节点的子树大小 size 和子树内距离和 dp。

第二次 DFS 通过换根公式计算其它节点的答案： ans[child] = ans[parent] - size[child] + (n - size[child])。

这样每条边只处理常数次，整体线性。

复杂度分析：

时间复杂度：O(n)，其中 n 表示节点数。

空间复杂度：O(n)。

import java.util.ArrayList;
import java.util.List;

class Solution {
    private List<Integer>[] graph;
    private int[] size;
    private int[] ans;
    private int n;

    public int[] sumOfDistancesInTree(int n, int[][] edges) {
        this.n = n;
        graph = new ArrayList[n];
        for (int i = 0; i < n; i++) {
            graph[i] = new ArrayList<>();
        }
        for (int[] e : edges) {
            graph[e[0]].add(e[1]);
            graph[e[1]].add(e[0]);
        }

        size = new int[n];
        ans = new int[n];

        dfs1(0, -1);
        dfs2(0, -1);

        return ans;
    }

    private void dfs1(int u, int parent) {
        size[u] = 1;
        for (int v : graph[u]) {
            if (v == parent) {
                continue;
            }

            dfs1(v, u);
            size[u] += size[v];
            ans[u] += ans[v] + size[v];
        }
    }

    private void dfs2(int u, int parent) {
        for (int v : graph[u]) {
            if (v == parent) {
                continue;
            }

            ans[v] = ans[u] - size[v] + (n - size[v]);
            dfs2(v, u);
        }
    }
}

func sumOfDistancesInTree(n int, edges [][]int) []int {
	graph := make([][]int, n)
	for _, e := range edges {
		u := e[0]
		v := e[1]
		graph[u] = append(graph[u], v)
		graph[v] = append(graph[v], u)
	}

	size := make([]int, n)
	ans := make([]int, n)

	var dfs1 func(u int, parent int)
	dfs1 = func(u int, parent int) {
		size[u] = 1
		for _, v := range graph[u] {
			if v == parent {
				continue
			}
			dfs1(v, u)
			size[u] += size[v]
			ans[u] += ans[v] + size[v]
		}
	}

	var dfs2 func(u int, parent int)
	dfs2 = func(u int, parent int) {
		for _, v := range graph[u] {
			if v == parent {
				continue
			}
			ans[v] = ans[u] - size[v] + (n - size[v])
			dfs2(v, u)
		}
	}

	dfs1(0, -1)
	dfs2(0, -1)

	return ans
}

相似题目

题目	难度	考察点
863. 二叉树中所有距离为 K 的结点	中等	树
1245. 树的直径	中等	树 DP
543. 二叉树的直径	简单	树 DP
979. 在二叉树中分配硬币	中等	树 DP
337. 打家劫舍 III	中等	树 DP

本文作者: HGNULB

本文链接: https://hgnulb.github.io/blog/2021/10883574