75.二叉树供暖问题|Marscode AI刷题

106 阅读7分钟

1.题目

问题描述

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

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


测试样例

样例1:

输入:nodes = [1, 1, 0, 1, 1]

输出:1

样例2:

输入:nodes = [1, 0, 1, 1, 0, 1, 0, 1, 0, 1, 0, 0, 1, 1]

输出:3

样例3:

输入:nodes = [1, 1, 0, 1, 1, 1, 1]

输出:2

2.思路

  • 构建二叉树

    • 输入是一个节点列表,其中 1 表示该节点存在,0 表示该节点不存在。我们需要根据这些节点信息,构建一棵二叉树。
    • 从列表中的第一个元素开始,根节点一定存在,然后逐级为每个节点分配左子节点和右子节点。使用队列来辅助构建树,逐个处理节点并为其分配子节点。
  • 深度优先搜索(DFS)

    • 我们的目标是计算每个子树中放置暖炉的最少数量。对于每个节点,可以有以下几种选择:

      • choose:在当前节点放置暖炉。
      • by_fa:当前节点的父节点已经放置暖炉,那么该节点可以不需要暖炉,或者根据情况放置暖炉。
      • by_children:当前节点的子节点放置暖炉时,当前节点可能也需要一个暖炉。
  • 状态转移

    • 对于每个节点,dfs 递归地计算其左右子树的状态值(即 choose, by_fa, by_children)。

    • 对于当前节点的左子树和右子树,分别考虑以下几种情况:

      • choose:当前节点放置暖炉时,左子树和右子树分别选择最优的方案(chooseby_fa)。
      • by_fa:当前节点的父节点已经放置暖炉,当前节点可以选择不放暖炉,或者根据子节点的状态决定是否需要放暖炉。
      • by_children:当前节点的子节点放置暖炉时,当前节点选择最优方案。
    • 然后返回这三种方案的最小值,作为当前节点的最优暖炉数量。

  • 返回结果

    • 从根节点开始递归计算,并返回根节点的最小暖炉数量。我们最终比较 chooseby_children 的值,选择最小值。

3.代码

from collections import deque # 导入双端队列,用于处理节点的层级遍历
from dataclasses import dataclass # 用于定义数据类
from typing import Optional # 用于表示可选类型(类型提示)

@dataclass
class TreeNode:
    # 定义二叉树节点的数据结构
    left: Optional['TreeNode'] = None
    right: Optional['TreeNode'] = None
def solution(nodes):
    # Please write your code here
    # 构建二叉树并找到至少需要多少个暖炉
    rt = TreeNode() # 创建根节点
    q1 = deque(nodes[1:]) # 从输入列表中获取节点信息(除去根节点)
    q2 = deque([rt]) # 创建一个队列,用于存储当前节点
    while q1: # 遍历节点
        u = q2.popleft() # 从队列中取出一个节点
        # 如果队列q1中还有元素,并且该元素为1,说明当前节点有左子节点
        if q1 and q1.popleft():
            v = TreeNode()
            u.left = v
            q2.append(v)
        # 如果队列q1中还有元素,并且该元素为1,说明当前节点有右子节点
        if q1 and q1.popleft():
            v = TreeNode()
            u.right = v
            q2.append(v)
    def dfs(node):
        # 深度优先遍历,并计算最少需要的暖炉数
        if node is None:
            return 10 ** 9, 0, 0 # 如果节点为空,返回很大的数,表示不需要暖炉
        # 递归遍历左子树和右子树
        l_choose, l_by_fa, l_by_children = dfs(node.left)
        r_choose, r_by_fa, r_by_children = dfs(node.right)
        # 选择在当前节点放暖炉时的情况
        choose = min(l_choose, l_by_fa) + min(r_choose, r_by_fa) + 1
        # 如果当前节点的父节点有暖炉,则计算父节点放暖炉的情况
        by_fa = min(l_choose, l_by_children) + min(r_choose, r_by_children)
        # 如果当前节点的子节点有暖炉,则计算子节点放暖炉的情况
        by_children = min(l_choose + r_by_children, r_choose + l_by_children, l_choose + r_choose)
        # 返回三种情况的结果
        return choose, by_fa, by_children
    # 调用dfs,返回的choose是从根节点开始的最少暖炉数
    choose, _, by_children = dfs(rt)

    # 返回最少需要的暖炉数
    return min(choose, by_children)

if __name__ == "__main__":
    #  You can add more test cases here
    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 )
  1. choose — 选择在当前节点放暖炉时的情况
  • 含义:这是当前节点本身选择放置暖炉时的最小暖炉数量。

  • 计算方式

    • l_choose 是左子树中选择放暖炉的最小数量。
    • l_by_fa 是左子树通过父节点来提供暖气的最小数量。
    • r_choose 是右子树中选择放暖炉的最小数量。
    • r_by_fa 是右子树通过父节点来提供暖气的最小数量。

    min(l_choose, l_by_fa)min(r_choose, r_by_fa) 分别表示在左子树和右子树中选择最优方案,不管是通过放暖炉在当前子树节点还是让父节点给子树提供暖气。然后再加上 1,表示当前节点本身也需要放一个暖炉来保证温暖。

  • 目的:为了让当前节点温暖,我们要从左子树和右子树的最优选择中选择最少的暖炉数量,并且加上当前节点的暖炉。


  1. by_fa — 如果当前节点的父节点有暖炉,则计算父节点放暖炉的情况
by_fa = min(l_choose, l_by_children) + min(r_choose, r_by_children)
  • 含义:这是通过父节点提供暖气的最小暖炉数量。

  • 计算方式

    • l_choose 是左子树中放置暖炉的最小数量。
    • l_by_children 是左子树的子节点提供暖气的最小数量。
    • r_choose 是右子树中放置暖炉的最小数量。
    • r_by_children 是右子树的子节点提供暖气的最小数量。

    min(l_choose, l_by_children) 表示左子树的最优选择,不论是放暖炉在左子树节点,还是通过子节点提供暖气,右侧同理。

  • 目的:如果当前节点没有直接放暖炉的情况,那么它可能需要通过父节点提供暖气。by_fa 计算的是这样的一种最优情况,即父节点给左右子树分别提供暖气的最少暖炉数量。


  1. by_children — 如果当前节点的子节点有暖炉,则计算子节点放暖炉的情况
by_children = min(l_choose + r_by_children, l_by_children + r_choose, l_choose + r_choose)
  • 含义:这是通过子节点提供暖气的最小暖炉数量。

  • 计算方式

    • l_choose 是左子树中放置暖炉的最小数量。
    • r_by_children 是右子树通过其子节点提供暖气的最小数量。
    • l_by_children 是左子树通过其子节点提供暖气的最小数量。
    • r_choose 是右子树中放置暖炉的最小数量。

    这里有三种情况:

    • l_choose + r_by_children:将暖炉放在左子树,并且右子树通过其子节点提供暖气。
    • l_by_children + r_choose:左子树通过其子节点提供暖气,右子树将暖炉放在其节点。
    • l_choose + r_choose:分别将暖炉放在左子树和右子树节点。
  • 目的:这是在考虑左右子树已经放置了暖炉的情况下,计算通过子节点(而非当前节点或父节点)来传递暖气的最优方式,得到最少的暖炉数量。