二叉树供暖问题（Python） | 豆包MarsCode AI刷题针对豆包MarsCode AI刷题中二叉树供暖问题，进

问题描述

已知村庄里每户人家作为一个节点，这些节点可以组成一个二叉树。我们可以在二叉树的节点上供应暖炉，每个暖炉可以为该节点的父节点、自身及其子节点带来温暖。给定一棵二叉树，求使整个村庄都暖和起来至少需要供应多少个暖炉？

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

输入：nodes数组，数组为按层遍历的二叉树各节点状况
输出：最小暖炉数量

问题分析

首先考虑各个节点可能存在的状态：

自身放置暖炉
已被覆盖 => 父节点放置暖炉
未覆盖

而根据题意我们需要使整个村庄都暖和起来（所有节点自身 or 父节点有暖炉），也就是说，我们需要根据子节点的状态（是否未覆盖），来决定当前节点的状态（是否放置暖炉）。

根据题目的这个特点，我们需要选择合适的二叉树遍历方式。

深度优先遍历：从未访问的节点出发，沿着一条路走到最深处，之后再回退到上一个节点，继续走别的路走到尽头......直到所有节点被访问。包括前序遍历，中序遍历，后序遍历。

前中后序是相对于根节点而言的
- 前序遍历：最先访问根节点，然后遍历左子树，最后遍历右子树
- 中序遍历：根结点排中间访问，先遍历左子树，然后访问根节点，最后遍历右子树。
- 后序遍历：根节点排最后访问，先遍历左子树，然后遍历右子树，最后访问根节点。
由于这题子节点和父节点的状态强相关，所以需要选用深度优先遍历。而根节点的左右子树决定了根节点的状态，所以需要在访问根节点之前，先遍历完左右子树。
广度优先遍历：按照层级顺序访问节点，从未访问的节点出发，先遍历节点的相邻节点，再依次遍历每个相邻节点的相邻节点。

题目输入的并不是二叉树本身，而只是按层遍历各个节点的数组

在solution函数中，还应构建辅助函数，用层次遍历的方式来逐层构建二叉树。

Python代码书写

定义节点数据结构

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

定义辅助函数build_tree

def build_tree(nodes):
        if not nodes:
            return None
        root = TreeNode() if nodes[0] == 1 else None
        queue = [root]
        i = 1
        while i < len(nodes):
            node = queue.pop(0)
            if node:
                node.left = TreeNode() if i < len(nodes) and nodes[i] == 1 else None
                queue.append(node.left)
                i += 1
                if i < len(nodes):
                    node.right = TreeNode() if nodes[i] == 1 else None
                    queue.append(node.right)
                    i += 1
        return root

初始化根节点：
- 如果 nodes[0] 为 1，则创建根节点 root，否则 root 为 None。
- 初始化队列 queue，将根节点 root 放入队列中。
逐层遍历：
- 从队列中取出队首节点 node。如果 node 不为 None，则继续处理其左右子节点。
处理左/右子节点：
- 如果 i 小于 nodes 的长度且 nodes[i] 为 1，则创建左/右子节点 node.left/node.right，并将其放入队列中。
- 将 i 增加 1。
最终返回构建好的二叉树的根节点 root

定义后序遍历函数

    def dfs(node):
        nonlocal heaters
        if not node:
            return 1  # 空节点认为是已覆盖的

        left = dfs(node.left)
        right = dfs(node.right)

        if left == 0 or right == 0:
            heaters += 1
            return 2  # 当前节点放置暖炉
        if left == 2 or right == 2:
            return 1  # 当前节点被覆盖

        return 0  # 当前节点未被覆盖

heaters为当前暖炉的数量

left和right的值分别表示该节点的子节点的状态

完整代码

def solution(nodes):
     # 定义树节点
    class TreeNode:
        def __init__(self):
            self.left = None
            self.right = None

    # 辅助函数，将列表转换为二叉树
    def build_tree(nodes):
        if not nodes:
            return None
        root = TreeNode() if nodes[0] == 1 else None
        queue = [root]
        i = 1
        while i < len(nodes):
            node = queue.pop(0)
            if node:
                node.left = TreeNode() if i < len(nodes) and nodes[i] == 1 else None
                queue.append(node.left)
                i += 1
                if i < len(nodes):
                    node.right = TreeNode() if nodes[i] == 1 else None
                    queue.append(node.right)
                    i += 1
        return root

    root = build_tree(nodes)
    heaters = 0

    def dfs(node):
        nonlocal heaters
        if not node:
            return 1  # 空节点认为是已覆盖的

        left = dfs(node.left)
        right = dfs(node.right)

        if left == 0 or right == 0:
            heaters += 1
            return 2  # 当前节点放置暖炉
        if left == 2 or right == 2:
            return 1  # 当前节点被覆盖

        return 0  # 当前节点未被覆盖

    # 如果根节点没有被覆盖，需再加一个暖炉
    if dfs(root) == 0:
        heaters += 1

    return heaters


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,0,1,1,0,0,1,0,1,1,0,0,1]) == 3 )