二叉树供暖问题(二叉树+后序遍历) | 豆包MarsCode AI刷题

99 阅读5分钟

一、题目分析

问题描述

天气越来越冷了,村子里有留守老人缺少照顾,会在这个冬天里挨冻,小R想运用自己的知识帮帮他们。已知村庄里每户人家作为一个节点,这些节点可以组成一个二叉树。我们可以在二叉树的节点上供应暖炉,每个暖炉可以为该节点的父节点、自身及其子节点带来温暖。给定一棵二叉树,求使整个村庄都暖和起来至少需要供应多少个暖炉?

本题按层遍历顺序描述二叉树的节点情况。值为 1,代表存在一个节点,值为 0,代表不存在该节点。

输入:给定一个二叉树的层序遍历表示,每个节点用 1 或 0 来表示是否存在。1 表示存在该节点,0 表示该节点不存在。

输出:我们需要输出最少需要多少个暖炉才能让整个村庄(即二叉树中的所有节点)都保持温暖。

二、核心思路

  1. 暖炉的影响范围:一个暖炉不仅仅给所在节点提供温暖,还能影响到父节点和子节点。所以暖炉的影响范围是整个节点的 "树状结构"。

  2. 节点状态的表示

    • 节点状态有 3 种:

      • 0 表示当前节点需要供暖,且放置了暖炉。
      • 1 表示当前节点已经被供暖,不需要再放置暖炉。
      • 2 表示当前节点未供暖,也没有放置暖炉。
  3. 递归遍历:为了判断整个二叉树是否需要供暖,采用后序遍历方式,从子节点到父节点依次判断:

    • 空节点:如果节点为空(即子节点为 None),默认该节点已经供暖。
    • 叶子节点:如果左右子节点都已经供暖,当前节点无需放置暖炉。
    • 部分未供暖的子节点:如果某个子节点未供暖,则需要在当前节点放置暖炉,从而为当前节点和未供暖的子节点提供温暖。
  4. 全局暖炉计数:通过递归方式统计最少需要多少个暖炉

三、解决方案

我们将使用后序遍历方式从下往上处理每个节点,并根据子节点的供暖情况决定是否放置暖炉。具体步骤如下:

  1. 构建二叉树:根据层序遍历的输入构建二叉树,树的结构会由节点的存在情况(1表示存在节点,0表示不存在节点)决定。
  2. 后序遍历并决策:递归地从每个子节点向上判断是否需要放置暖炉,并对全局暖炉计数进行更新。
  3. 计算结果:在遍历结束后,返回最少需要的暖炉数。

四、代码实现

from collections import deque

# 定义树节点
class Node:
    def __init__(self, data):
        self.data = data
        self.left = None
        self.right = None

# 记录暖炉个数
result = 0

# 解决函数
def solution(nodes):
    global result
    result = 0  # 每次调用时初始化为0,防止沿用之前的结果
    
    # 根据层序遍历的数组构造二叉树
    root = build_tree(nodes)
    
    # 判断根节点是否已经供暖
    if postorder_traversal(root) == 2:
        result += 1
    
    return result

# 后序遍历
# 每个节点有三种情况:0 -> 安装暖炉  1 -> 已供暖  2 -> 未供暖
def postorder_traversal(root):
    # 空节点的情况为被供暖
    if root is None:
        return 1
    
    # 左子树
    left = postorder_traversal(root.left)
    # 右子树
    right = postorder_traversal(root.right)
    
    # 根节点判断
    # 如果两个叶子节点都已供暖
    if left == 1 and right == 1:
        # 返回当前节点未供暖
        return 2
    
    # 如果两个叶子节点有一个未供暖
    if left == 2 or right == 2:
        # 安装暖炉给叶子节点供暖
        global result
        result += 1
        return 0
    
    # 如果两个叶子节点其中一个安装了暖炉
    if left == 0 or right == 0:
        return 1
    
    # 实际上不会走到这一步,上面已经将所有情况都考虑了
    return 1

# 根据层序遍历的数组构建二叉树
def build_tree(nodes):
    if not nodes:
        return None
    
    # 构建根节点
    root = Node(nodes[0])
    # 创建一个队列用于存储节点
    queue = deque([root])
    # 从数组的第二个元素开始
    i = 1
    while queue and i < len(nodes):
        current = queue.popleft()
        
        # 处理左子节点
        if i < len(nodes) and nodes[i] == 1:
            current.left = Node(nodes[i])
            queue.append(current.left)
        i += 1
        
        # 处理右子节点
        if i < len(nodes) and nodes[i] == 1:
            current.right = Node(nodes[i])
            queue.append(current.right)
        i += 1
    
    return root

# 测试
if __name__ == "__main__":
  
    print(solution([1, 1, 0, 1, 1]) == 1)  # 只需要对[0]做操作
    print(solution([1, 0, 1, 1, 0, 1, 0, 1, 0, 1, 0, 0, 1, 1]) == 3)  # 有三个连续0段
    print(solution([1, 1, 0, 0, 1, 1, 0, 0, 1, 0, 1, 1, 0, 0, 1]) == 3)  # 有三个连续0段

代码解释

  1. Node 类:定义了二叉树的节点结构,每个节点包含一个数据域 data,以及左右子节点 leftright
  2. build_tree 函数:根据给定的层序遍历输入数组,构造一棵二叉树。通过队列 queue 逐层遍历并创建每个节点,左右子节点依次被赋值。
  3. postorder_traversal 函数:递归的后序遍历方法,判断每个节点是否已经温暖,并决定是否需要在当前节点安装暖炉。遍历过程中需要判断当前节点的左右子树是否已经被供暖,并通过返回值更新全局暖炉计数 result
  4. solution 函数:这是主函数,首先构建二叉树,然后调用 postorder_traversal 函数来判断并计算最少需要多少个暖炉

五、总结

本题的核心问题是如何通过后序遍历判断并计算出最少的暖炉数量,确保每个节点都被温暖到。通过合理的递归和状态判断,我们能够高效地解决问题。

一开始我对该题目与过程感到非常苦恼,因为二叉树问题本身就较为复杂,而且涉及到暖炉的放置规则让我有些迷惑。通过AI的思路提示,我能够逐渐理清思路,明白了如何通过后序遍历从叶子节点往上判断每个节点是否需要放置暖炉,以及如何根据子节点的状态来决定父节点的操作。通过这个过程,我逐步掌握了问题的解法,感到很有成就感。