题目解析:二叉树供暖问题 | 豆包MarsCode AI刷题

4 阅读6分钟

二叉树供暖问题

一、问题分析

村庄里的留守老人面临寒冷的问题,需要通过安装暖炉为他们提供温暖。村庄的住户分布可以用一个二叉树来描述,其中每个节点代表一户人家。如果在某个节点安装了暖炉,该暖炉可以同时温暖当前节点、其父节点及其所有子节点。这意味着一个暖炉的覆盖范围不局限于单个节点,而是可以涵盖多个相关节点,因此问题的关键在于如何合理地分配暖炉,使得村庄中所有住户节点都能够被覆盖。

给定的输入以层序遍历的方式描述二叉树的节点分布,其中值为 1 表示节点存在,值为 0 表示节点不存在。目标是根据这一结构判断需要安装暖炉的最少数量,以最优化的方式满足覆盖所有节点的要求。需要注意的是,输入中的空节点需要跳过处理,因此二叉树的构建及递归计算需要结合二叉树的层序特性。

本问题本质上是一个二叉树节点覆盖的优化问题,需要基于节点的状态进行合理推导和规划。通过分析可以发现,每个节点有三种状态:未覆盖、被覆盖但未安装暖炉、已安装暖炉。合理地利用这些状态并设计递归规则,将帮助我们从整体上求解最优方案。

二、思路解析

解决该问题需要采用一种从下到上的递归贪心策略。二叉树中的叶子节点是分析的起点,如果一个叶子节点未被覆盖,则其父节点必须安装暖炉。同理,如果一个节点的子节点中存在未被覆盖的情况,那么当前节点也需要安装暖炉。这种策略能够确保覆盖范围尽可能广,避免多余的暖炉安装。同时,如果一个节点的子节点中已经安装了暖炉,那么当前节点就可以被认为是被覆盖的,不需要额外安装暖炉。

在实现过程中,首先需要将输入的层序遍历数组转换为实际的二叉树结构,这一步需要通过队列实现。此外,在递归计算过程中,每个节点的状态都需要根据其左右子节点的状态动态更新。如果左右子节点均被覆盖且没有暖炉,那么当前节点状态为未覆盖;如果任意子节点未覆盖,则当前节点需要安装暖炉;如果子节点中存在暖炉,则当前节点已被覆盖但未安装暖炉。根节点的状态在递归结束后需要单独检查,如果未被覆盖,则需要增加一个暖炉。

最终,通过这种递归和状态转移规则,能够有效地计算出安装暖炉的最少数量。整个过程的时间复杂度为 O(n),其中 n 是节点总数,空间复杂度为递归深度 O(h),其中 h 是二叉树的高度。

三、代码详解

class TreeNode:
    def __init__(self, val=0):
        self.val = val
        self.left = None
        self.right = None

def build_tree(nodes):
    """
    根据层序遍历构建二叉树
    """
    if not nodes or nodes[0] == 0:
        return None
    root = TreeNode(nodes[0])
    queue = [root]
    i = 1
    while i < len(nodes):
        current = queue.pop(0)
        if current:
            if i < len(nodes) and nodes[i] != 0:
                current.left = TreeNode(nodes[i])
                queue.append(current.left)
            i += 1
            if i < len(nodes) and nodes[i] != 0:
                current.right = TreeNode(nodes[i])
                queue.append(current.right)
            i += 1
    return root

def min_heaters(root):
    """
    计算需要的最少暖炉数量
    """
    def dfs(node):
        if not node:
            return 1  # 空节点视为已覆盖
        left = dfs(node.left)
        right = dfs(node.right)
        if left == 0 or right == 0:
            # 如果子节点未被覆盖,需要放置暖炉
            nonlocal heaters
            heaters += 1
            return 2
        if left == 2 or right == 2:
            # 如果子节点有暖炉,则当前节点被覆盖
            return 1
        return 0  # 当前节点未被覆盖

    heaters = 0
    if dfs(root) == 0:  # 如果根节点未被覆盖,则需要放置一个暖炉
        heaters += 1
    return heaters

def solution(nodes):
    root = build_tree(nodes)
    return min_heaters(root)

if __name__ == "__main__":
    print(solution([1, 1, 0, 1, 1]) == 1)
    print(solution([1, 0, 1, 1, 0, 1, 0, 1, 0, 1, 0, 0, 1, 1]) == 3)
    print(solution([1, 1, 0, 1, 1, 1, 1]) == 2)

代码的逻辑实现分为三个主要部分:二叉树的构建、暖炉数量的递归计算以及整体的解决方案。首先,通过 build_tree 函数将输入的层序遍历数组转化为实际的二叉树结构,这一步使用了队列以维护节点的逐层关系,从根节点开始依次为其左右子节点赋值,遇到值为 0 的位置跳过,以构建完整的二叉树结构。随后,min_heaters 函数负责计算所需的最少暖炉数量,其核心在于递归函数 dfs。在 dfs 中,通过深度优先遍历从叶子节点向根节点反向更新每个节点的状态。根据子节点的状态决定当前节点的处理方式:如果任一子节点未被覆盖,则需要在当前节点安装暖炉;如果子节点中存在暖炉,则当前节点已被覆盖,无需再放置暖炉;否则,当前节点未被覆盖。递归结束后,还需检查根节点的状态,若根节点未被覆盖,则额外放置一个暖炉以确保覆盖范围完整。最后,solution 函数负责整合前两部分,将输入数组构建为二叉树后调用 min_heaters 返回结果,整体逻辑清晰且高效。

四、总结

本问题通过贪心策略结合递归的方法,高效地解决了村庄暖炉分布优化的问题。核心思路在于利用二叉树的结构特点,从叶子节点向根节点递归计算覆盖状态,以确定最优的暖炉放置位置。这种自底向上的递归方式保证了每个节点都在最小的暖炉数量下被有效覆盖,同时避免了不必要的重复计算。通过三种状态的设计(未覆盖、已覆盖但无暖炉、放置暖炉),问题被自然地分解为局部最优解并最终汇总为全局最优解。此外,通过深度优先搜索,递归调用不仅逻辑简单,还能有效控制时间复杂度,使整体算法的效率能够满足大规模输入的需求。

整个方案在构建二叉树、递归状态转移和最终结果输出的流程上实现了清晰的分工和高效的计算。代码中使用了灵活的队列构建层序二叉树,并结合动态状态更新,充分利用了二叉树的结构特性。最终,算法的时间复杂度为 O(n),空间复杂度由递归深度决定为 O(h)。这种方法不仅适用于当前问题,还可以推广到其他类似的覆盖优化问题,为解决此类问题提供了通用的设计思路和实现方式。