当前位置: 面试刷题>> 树中距离的总和 (经典算法题500道)


题目描述补充

给定一棵无向树(即图论中的树结构,无环且连通),其中节点从1到N编号。树的边有权重,表示节点之间的距离。现要求计算树中所有节点对之间的距离总和。

示例

假设有以下的树结构,边和权重如下:

    1
   / \
  2   3
 / \   \
4   5   6

边和权重为:

  • (1, 2) = 1
  • (1, 3) = 2
  • (2, 4) = 3
  • (2, 5) = 4
  • (3, 6) = 5

我们需要计算这棵树中所有节点对之间的距离总和。

解题思路

一个有效的解决方案是使用树的遍历和组合计数的思想。基本思路是对于树中的每个节点,我们可以计算出以该节点为根的子树中所有节点到其他非子树节点的距离之和。这可以通过两次DFS(深度优先搜索)实现:第一次DFS计算每个子树的大小和到根的距离之和,第二次DFS利用第一次DFS的信息来计算总距离。

PHP 示例代码

<?php
class TreeNode {
    public $val;
    public $children;

    function __construct($val = 0, $children = []) {
        $this->val = $val;
        $this->children = $children;
    }
}

function sumOfDistancesInTree($root) {
    $totalNodes = 0;
    $subtreeSizes = [];
    $subtreeSums = [];

    // 第一次DFS
    function dfs1($node, &$parent = null) use (&$totalNodes, &$subtreeSizes, &$subtreeSums) {
        $totalNodes++;
        $subtreeSizes[$node->val] = 1;
        $subtreeSums[$node->val] = 0;

        foreach ($node->children as $child) {
            dfs1($child, $node);
            $subtreeSizes[$node->val] += $subtreeSizes[$child->val];
            $subtreeSums[$node->val] += $subtreeSums[$child->val] + $subtreeSizes[$child->val];
        }

        if ($parent !== null) {
            $subtreeSums[$parent->val] += $totalNodes - $subtreeSizes[$node->val];
        }
    }

    dfs1($root);

    $totalDistance = 0;

    // 第二次DFS,利用第一次DFS的结果计算总距离
    function dfs2($node, $parentDist = 0) use (&$totalDistance, &$subtreeSizes, &$subtreeSums, $totalNodes) {
        $totalDistance += $subtreeSums[$node->val] + ($totalNodes - $subtreeSizes[$node->val]) * $parentDist;

        foreach ($node->children as $child) {
            $newParentDist = $parentDist + $totalNodes - $subtreeSizes[$child->val];
            dfs2($child, $newParentDist);
        }
    }

    dfs2($root);

    return $totalDistance;
}

// 示例用法
$root = new TreeNode(1, [
    new TreeNode(2, [
        new TreeNode(4),
        new TreeNode(5)
    ]),
    new TreeNode(3, [
        new TreeNode(6)
    ])
]);

echo sumOfDistancesInTree($root);  // 输出计算结果
?>

注意

  • 上面的PHP代码是一个概念验证示例,具体实现可能需要根据实际的数据结构和环境进行调整。
  • PHP代码中的TreeNode类用于表示树节点,并包含节点值和子节点列表。
  • 两次DFS分别用于收集和计算距离信息。
  • 码小课网站中有更多关于树结构、DFS、图论等相关内容的分享,可以进一步学习和探索。
推荐面试题