树上的DP - Set-3(N进制树的直径)

在树结构上的动态规划被称为“树上DP”。其中，N进制树的直径则是一道经典的树上DP问题。在本文中，我们将介绍N进制树的直径问题以及相关的算法思路和代码实现。

问题描述

给定一棵由N进制数字构成的树，每个节点上都有一个N进制数字。树的直径定义为所有节点距离的最大值。

解题思路

树的直径问题解法通常有两种算法方法：DFS和动态规划。这里我们介绍动态规划的算法思路。

在本题中，我们需要用到两个DP数组：$dp1$和$dp2$。其中，$dp1[i]$用于表示节点$i$往子树方向的路径长度最大值；$dp2[i]$用于表示节点$i$往父节点方向的路径长度最大值。

我们首先需要预处理出每个节点的父节点，以及从根节点开始到每个节点的路径长度。设$fa[i]$表示节点$i$的父节点，$dis[i]$表示从根节点开始到节点$i$的路径长度。

DP转移方程如下所示：

$$ dp1[i]=\max_{(j,i)\in E}{dp1[j] + dist[j][i]} $$

$$ dp2[i]=\max_{fa[i]\neq 0}{dp2[fa[i]] + dist[i][fa[i]]} $$

其中，$dist[i][j]$表示节点$i$和节点$j$之间的距离，等于节点$i$到根节点的距离加上节点$j$到根节点的距离减去两者之间的公共部分。由于本题中每个节点是N进制数字，因此我们需要使用一个函数$dis(x)$来计算节点$x$到根节点的距离。

计算树的直径时，我们需要枚举每一个节点$i$，计算$dp1[i]+dp2[i]$的最大值。

代码实现

下面是本题的参考代码实现:

# 树的直径问题
# N进制树的直径问题
from typing import List


def dis(x: int) -> List[int]:
    """
    计算节点x到根节点的距离
    """
    d = []
    while x:
        d.append(x % n)
        x //= n
    for i in range(len(d), h):
        d.append(0)
    return d[::-1]


def solve():
    """
    计算树的直径
    """
    # 初始化
    for u in range(1, num + 1):
        dp1[u] = dp2[u] = 0
    # 计算dp1
    for i in range(num, 0, -1):
        for j in edges[i]:
            if j > i:
                continue
            d = dis(j)
            if i != j:
                dp1[i] = max(dp1[i], dp1[j] + sum(d))
            for k in range(h):
                dp1[i] = max(dp1[i], dp2[j] + sum(d[k:]))
    # 计算dp2
    for i in range(1, num + 1):
        if fa[i]:
            d = dis(i)
            for k in range(h):
                dp2[i] = max(dp2[i], dp2[fa[i]] + sum(d[k:]))
    # 计算直径
    ans = 0
    for i in range(1, num + 1):
        ans = max(ans, dp1[i] + dp2[i])
    return ans


# 主函数
if __name__ == '__main__':
    n = 3
    h = 20
    num = 0
    fa = [0] * (1 << h)
    edges = [[] for _ in range(1 << h)]
    dp1 = [0] * (1 << h)
    dp2 = [0] * (1 << h)
    # 读入数据
    with open('input.txt', 'r') as f:
        num = int(f.readline())
        for i in range(1, num + 1):
            x, y = map(int, f.readline().split())
            fa[i] = x
            edges[x].append(i)
    # 计算答案
    ans = solve()
    # 输出结果
    with open('output.txt', 'w') as f:
        f.write(str(ans) + '\n')

以上程序采用了Python编程语言实现。其中，$n$表示N进制数的底数（如$n=3$表示三进制数）；$h$表示树的高度，取值范围为$[1,20]$；$num$表示树的节点数量；$fa$表示每个节点的父节点；$edges$表示树的边；$dp1$和$dp2$分别为前文中提到的DP数组。

总结

本文中，我们介绍了N进制树的直径问题，并以动态规划作为算法解题方法。DP时使用了两个DP数组$dp1$和$dp2$，分别用于记录子树和父节点方向的路径最大值。计算树的直径时，可枚举每个节点$i$，计算$dp1[i]+dp2[i]$的最大值即可。