华为华为AI岗实习面试算法题

Prompt上下文信息精简

在一次 Prompt 工程实践中，我们把一段 token 序列抽象成一棵二叉树。树中每个结点都有一个整数权值（可正可负，也可能为 0）。请在这棵树中选出一棵"价值最大"的子树，并把这棵子树按"完全二叉树的层序数组"形式输出。

子树的价值定义为它所包含的所有结点权值之和。

允许对某个结点"剪掉"对总和贡献为负的整棵子树（即可以只要左子树、或只要右子树、或两者都要；被剪掉的位置在输出中以 null 占位）。

输入是一棵"用层序数组表示的完全二叉树"，缺失位置用 null 占位；输出也使用相同规则表示挑选出的那棵最优子树，并且去除末尾多余的尾部 null。

思路：

这是一个树形 DP（Dynamic Programming on Tree）求最大子树和的问题

1，首先根据DFS（递归）的方式和DP的思想，来找到每个节点的最大子树和，用value数组表示，

2，value最大的节点即为目标根节点

3，还原出最大子树的输出结构（输出的最大子树需要重新分配序号），使用递归的方式，分别从左子树和右子树来构造

4，输出格式化

max(key=function) 根据函数的规则来比较大小

代码实现：

import sys
sys.setrecursionlimit(10**7)

# 读取一行，如：[5,-1,3,null,null,4,7]
line = sys.stdin.read().strip()
line = line.strip()[1:-1]  # 去掉方括号

tokens = [x.strip() for x in line.split(',')] if line else []
arr = []
for t in tokens:
    if t.lower() == 'null' or t == '':
        arr.append(None)
    else:
        arr.append(int(t))

n = len(arr)
if n == 0:
    print("[]")
    sys.exit(0)

# 存储 DP 结果
value = [float("-inf")] * n
keepL = [False] * n
keepR = [False] * n

def dfs(i):
    if i >= n or arr[i] is None:
        return float("-inf")

    left = dfs(2 * i + 1)
    right = dfs(2 * i + 2)

    best_left = max(left, 0)
    best_right = max(right, 0)

    keepL[i] = (left > 0)
    keepR[i] = (right > 0)

    value[i] = arr[i] + best_left + best_right
    return value[i]

dfs(0)

best_root = max(range(n), key=lambda i: value[i])

# 构造新子树（使用字典存储稀疏完全二叉树）
newtree = {}

def build(i, pos):
    if i >= n or arr[i] is None:
        newtree[pos] = None
        return

    newtree[pos] = arr[i]

    # 左
    if keepL[i]:
        build(2*i+1, 2*pos+1)
    else:
        newtree[2*pos+1] = None

    # 右
    if keepR[i]:
        build(2*i+2, 2*pos+2)
    else:
        newtree[2*pos+2] = None

build(best_root, 0)

# 构造新子树（使用字典存储稀疏完全二叉树）
# newtree = {} 假设已经构造好了
# 转为数组形式
max_pos = max(newtree.keys()) # 设置最大长度
out = [None] * (max_pos + 1)
for k, v in newtree.items():
    out[k] = v

# 去除末尾 null
while out and out[-1] is None:
    out.pop()

# 输出格式：[x,y,z]
res = []
for v in out:
    if v is None:
        res.append("null")
    else:
        res.append(str(v))

print("[" + ",".join(res) + "]")