python coding with ChatGPT 专题1

python coding with ChatGPT 专题1 | 树的直径

文章目录

定义
题目
- 特点
树的表示
- 字典存储邻接表
- TreeNode类
深度优先（两次DFS法）
动态规划（树形DP）
- 优势
相似题目
参考资料

定义

树上任意两节点之间最长的简单路径即为树的「直径」。

题目

给定一棵树，求出这棵树的直径，即树上最远两点的距离。
包含n个结点，n-1条边的连通图称为树。
示例1的树如下图所示。其中4到5之间的路径最长，是树的直径，距离为5+2+4=11

在这里插入图片描述

输入：

edge_list = [[0,1],[1,5],[1,2],[2,3],[2,4]]
edge_value = [3,4,2,1,5]
n = 6

特点

一个没有固定根结点的树称为无根树（unrooted tree）。无根树有几种等价的形式化定义：

有 [n] 个结点， [n-1] 条边的连通无向图
无向无环的连通图
任意两个结点之间有且仅有一条简单路径的无向图
任何边均为桥的连通图
没有圈，且在任意不同两点间添加一条边之后所得图含唯一的一个圈的图

树的表示

字典存储邻接表

def get_tree(edge_list, edge_value, n):
    tree = {}
    for i in range(n):
        tree[i] = []
    for x, y in zip(edge_list, edge_value):
        tree[x[0]].append((x[1], y))
        tree[x[1]].append((x[0], y))
    return tree

TreeNode类

class TreeNode:
	def __init__(val):
		self.val = val
		self.neighbors = []
	
	def add_neighbor(self, node, weight):
		self.neighbors.append((node, weight))

def build_tree(edge_list, edge_value, n):
	nodes = [TreeNode(i) for i in range(n)]
	for (start, end), weight in zip(edge_list, edge_value):
		nodes[start].add_neighbor(end, y)
		nodes[end].add_neighbor(start, y)
	return nodes[0]

深度优先（两次DFS法）

解决这个问题的一个有效方法是利用深度优先搜索（DFS）。这个方法可以分为两步来完成：

从任意一个结点出发，执行一次深度优先搜索（DFS），找到距离它最远的结点。
以步骤1中找到的结点作为起点，再次执行深度优先搜索（DFS），找到距离它最远的结点的距离，这个距离即为树的直径。

为什么这个方法有效？第一次DFS帮助我们定位到树的一端，然后第二次DFS从这一端出发能够找到最远的另一端，这就保证了我们计算出的距离是最长的路径。

在这里插入图片描述

dfs逻辑：

# simple dfs for dict
def simple_dfs(current_node, parent, tree)
	print(current_node)
	for node, _ in tree[current_node):
		if node != parent:
			simple_dfs(node)

# simple dfs for treenode
def simple_dfs(current_node, parent)
	print(current_node)
	for node _ in current_node.neighbors:
		if node != parent:
			simple_dfs(node)

解题代码：

方法一：字典形式

# dict
def get_tree(edge_list, edge_value, n):
    tree = {}
    for i in range(n):
        tree[i] = []
    for x, y in zip(edge_list, edge_value):
        tree[x[0]].append((x[1], y))
        tree[x[1]].append((x[0], y))
    return tree
    
def dfs(current_node, parent, tree, distance):
	farthest_node = current_node
	max_distance = distance
	for node, weight in tree[current_node]:
		if node != parent:
			this_farthest_node, this_max_distance = dfs(node, current_node, tree, distance+weight)
			if this_max_distance > max_distance:
				max_distance = this_max_distance
				farthest_node = this_farthest_node
	return farthest_node, max_distance

def treeDiameter(edge_list, edge_value, n):
	tree = get_tree(edge_list, edge_value, n)
	farthest_node, _ = dfs(0, None, tree, 0)
	_, max_distance = dfs(farthest_node, None, tree, 0)
	return max_distance

方法二：TreeNode形式

class TreeNode:
    def __init__(self, val):
        self.val = val
        self.neighbors = []

    def add_neighbor(self, node, weight):
        self.neighbors.append((node, weight))


def build_tree(edge_list, edge_value, n):
    nodes = [TreeNode(i) for i in range(n)]
    for (start, end), weight in zip(edge_list, edge_value):
        nodes[start].add_neighbor(nodes[end], weight)
        nodes[end].add_neighbor(nodes[start], weight)
    return nodes[0]

def dfs(current_node, parent, distance):
    farthest_node = current_node
    max_distance = distance
    for node,weight in current_node.neighbors:
        if node != parent:
            this_farthest_node, this_max_distance = dfs(node, current_node, distance+weight)
            if this_max_distance > max_distance:
                farthest_node = this_farthest_node
                max_distance = this_max_distance
    return farthest_node, max_distance

def treeDiameter(edge_list, edge_value, n):
    current_node = build_tree(edge_list, edge_value, n)
    farthest_node, _ = dfs(current_node, None, 0)
    _, max_distance = dfs(farthest_node, None, 0)
    return max_distance

动态规划（树形DP）

定义一个DP状态来表示从当前节点出发能够达到的最长路径长度
这个方法的关键在于，我们在遍历树的过程中，考虑通过这个节点的最长路径(也就是这个节点作为路径中间点，连接两个最长的子路径)。
在这里插入图片描述

def calculate_diameter(node, parent, diameter):
    max_depth1, max_depth2 = 0, 0  # 存储当前节点的最大和次大深度

    for neighbor, weight in node.neighbors:
        if neighbor == parent:
            continue

        # 获取到该邻接节点的最大深度
        depth = calculate_diameter(neighbor, node, diameter) + weight

        # 更新最大和次大深度
        if depth > max_depth1:
            max_depth2, max_depth1 = max_depth1, depth
        elif depth > max_depth2:
            max_depth2 = depth

    # 更新通过当前节点的最大直径
    diameter[0] = max(diameter[0], max_depth1 + max_depth2)

    # 返回到当前节点的最大深度
    return max_depth1

def treeDiameter(edge_list, edge_value, n):
    root = build_tree(edge_list, edge_value, n)
    diameter = [0]  # 使用列表作为引用传递，以便在递归中更新
    calculate_diameter(root, None, diameter)
    return diameter[0]