DP 树 | Set-3（N 叉树的直径）

本文将介绍如何用动态规划解决 N 叉树的直径问题，包括问题的定义、解法思路、伪代码和示例代码。

问题定义

N 叉树的直径定义为树中两个节点之间的最长路径。换句话说，它是树中所有路径长度中最大的一条。

解法思路

N 叉树的直径可以通过动态规划来解决。假设一棵树的根节点为 root，我们可以通过以下步骤来求出该树的直径：

1. 对于树中的每个节点 u，计算它到根节点的最长路径 d1(u) 和次长路径 d2(u)，其中 d1(u) 为 u 到根节点的最长路径长度，d2(u) 为 u 到根节点的次长路径长度。
2. 对于树中的每个节点 u，计算以它为路径中的一个端点的直径 d3(u)，其中 d3(u) 为 u 到它的某个子节点 v 的路径长度加上 v 的以它为路径中的一个端点的直径。
3. 对于树中的每个节点 u，将 d3(u)、d1(u) 和 d2(u) 中的最大值作为以 u 为路径中的一个端点的直径。

最终的直径为树中所有路径长度中的最大值，即树中节点的 d3(u) 中的最大值。

伪代码

下面是求解 N 叉树直径的伪代码：

# 计算节点到根节点的最长和次长距离
def calc_distance(u, parent):
    d1, d2 = 0, 0
    for v in u.children:
        if v != parent:
            d = calc_distance(v, u)
            if d > d1:
                d1, d2 = d, d1
            elif d > d2:
                d2 = d
    u.d1, u.d2 = d1 + 1, d2 + 1

# 计算以节点为路径中的一个端点的直径长度
def calc_diameter(u, parent):
    u.d3 = 0
    for v in u.children:
        if v != parent:
            calc_diameter(v, u)
            d = u.d3 + v.d1 + 1
            if d > u.d3:
                u.d3 = d
    if u.d3 > u.diameter:
        u.diameter = u.d3
    if u.d1 + u.d2 > u.diameter:
        u.diameter = u.d1 + u.d2

# 求解树的直径
def calc_tree_diameter(root):
    calc_distance(root, None)
    calc_diameter(root, None)
    return root.diameter

示例代码

下面是一个示例程序，它输入一棵 N 叉树并输出它的直径：

class TreeNode:
    def __init__(self, value):
        self.value = value
        self.children = []
        self.d1, self.d2, self.d3, self.diameter = 0, 0, 0, 0

def build_tree(n, edges):
    nodes = [TreeNode(i) for i in range(n)]
    for u, v in edges:
        nodes[u].children.append(nodes[v])
        nodes[v].children.append(nodes[u])
    return nodes[0]

def test():
    edges = [(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (4, 6), (6, 7)]
    root = build_tree(8, edges)
    assert calc_tree_diameter(root) == 5

在上面的示例程序中，我们输入一棵 8 个节点的 N 叉树，其中 (0, 1) 表示节点 0 和节点 1 之间有一条边，输出该树的直径，预期输出为 5。