二叉树供暖问题 ｜ 豆包MarsCode AI 刷题

问题描述

天气越来越冷了，村子里有留守老人缺少照顾，会在这个冬天里挨冻，小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

分析

本题实际上是在要求完成两件任务：

• 将输入nodes转化成一棵二叉树T（只需关注nodes中存在的节点）；
• 找到在T上成本最低的暖炉安装策略。

将输入nodes转化成二叉树T

nodes是T的层序遍历；对于不存在的节点（0），其不存在的左右孩子节点皆被省略，不会在nodes的后续部分出现。

将层序遍历序列nodes转化为二叉树T的经典方法是：

1. 借助一个队列queue，queue最初只有nodes的根节点；
2. 当queue不为空时，出队得node，否则前往第4步。
3. 若node的左子节点left存在，则将node与left的关系加入T，left加入queue；若node的右子节点right存在，则将node与right的关系加入T，right加入queue；回到第2步。
4. 返回T；

许多题解会使用记录节点信息（如节点的父节点、左右子节点以及节点本身的值）的结构体或类来建构二叉树。这种方法的好处是查找父子节点的时间复杂度为$O(1)$，缺点是有一定空间开销。

笔者使用节点索引值的序列（题解中的indices）来表示二叉树。其原理是，在完全二叉树中，如果根节点的索引是1，那么索引为$n$的节点的左右孩子分别具有索引$2n$和$2n+1$。这种方法的空间需求小，而查找父子节点的时间复杂度为$O(\log n)$（indices是单调递增的，可使用二分查找）。

找到在T上成本最低的暖炉安装策略

最优策略可以通过贪心算法找到。每个节点有三种“亮度”状态2、1和0， 分别对应已安装暖炉、相邻节点已安装暖炉和未安装暖炉且相邻节点未安装暖炉。 每当遇到2，暖炉安装数就加一。从根节点开始，自顶向下递归进行：

• 不存在的节点为1；
• 若左右子节点皆为1，则当前节点为0；（叶子节点就是这种情况。为了避免浪费，叶子节点不能安装暖炉，而应由其父节点安装。）
• 如果左右子节点有一个为0，那么当前节点必须为2；
• 否则，若左右子节点有一个为2，则当前节点为1；

题解

from bisect import bisect_right
from typing import Literal


def solution(nodes: list[int]) -> int:
    ans = 0
    if not nodes or nodes[0] == 0:
        return ans
    indices: list[int] = [1]  # 所有存在的节点的索引
    parent, parents = 1, []
    i, size = 1, len(nodes)
    while i < size:
        # 将左右子节点加入队列`parents`
        if nodes[i]:
            parents.append(index := 2 * parent)
            indices.append(index)
        i += 1
        if i == size:
            break
        if nodes[i]:
            parents.append(index := 2 * parent + 1)
            indices.append(index)
        i += 1
        if parents:
            parent = parents.pop(0)
        else:
            break

    def inspect(index: int) -> Literal[0, 1, 2]:
        """
        Returns
        -------
            `0`: not covered;
            `1`: covered;
            `2`: installed.
        """
        nonlocal ans
        if indices[bisect_right(indices, index) - 1] != index:
            return 1
        left, right = inspect(double := 2 * index), inspect(double + 1)
        if left == right == 1:
            return 0
        if 0 in (left, right):
            ans += 1
            return 2
        if 2 in (left, right):
            return 1
        return 0

    if inspect(1) == 0:
        ans += 1

    return ans