CCF-CSP认证考试 202312-3 树上搜索：从模拟到优化的算法精讲-平芜编程栈

1. 题目背景与核心问题

这道CCF-CSP认证考试的第三题，描述了一个名词分类系统的场景。系统通过树形结构组织名词类别，每个类别都有权重值表示名词属于该类的可能性。题目要求实现一个二分策略，通过尽可能少的提问次数确定名词的具体分类。

实际解题时，我们需要处理的关键点包括：

树形结构的表示与遍历
动态计算子树权重和
每次迭代选择最优提问节点
根据用户回答更新搜索范围

我在第一次尝试时，直接用邻接表存储树结构，然后每次暴力计算子树权重和。这样虽然能通过样例，但当树规模达到2000节点时，时间复杂度会达到O(n²m)，导致最后两个测试点超时。

2. 基础解法：暴力模拟思路

先来看最直接的模拟解法。我们需要实现题目描述的四个步骤：

def brute_force_solution(): # 读取输入 n, m = map(int, input().split()) weights = list(map(int, input().split())) parents = list(map(int, input().split())) # 构建树结构 tree = [[] for _ in range(n+1)] for i in range(2, n+1): tree[parents[i-2]].append(i) # 处理每个测试用例 for _ in range(m): target = int(input()) current_nodes = set(range(1, n+1)) result = [] while len(current_nodes) > 1: # 计算每个节点的w_delta min_delta = float('inf') best_node = None for node in current_nodes: # 计算子树权重和 subtree_sum = 0 stack = [node] while stack: u = stack.pop() if u in current_nodes: subtree_sum += weights[u-1] stack.extend(tree[u]) total_sum = sum(weights[u-1] for u in current_nodes) delta = abs(2*subtree_sum - total_sum) if delta < min_delta or (delta == min_delta and node < best_node): min_delta = delta best_node = node result.append(str(best_node)) # 模拟用户回答 stack = [best_node] subtree_nodes = set() while stack: u = stack.pop() subtree_nodes.add(u) stack.extend(tree[u]) if target in subtree_nodes: current_nodes = subtree_nodes else: current_nodes -= subtree_nodes print(' '.join(result))

这个解法的问题在于，每次选择提问节点时都要重新计算所有节点的子树权重和，导致时间复杂度过高。在实际测试中，当n=2000时，运行时间会超过1秒的限制。

3. 优化思路：预处理与动态维护

观察发现，每次迭代只是从树中移除部分节点，而树的基本结构并未改变。我们可以通过以下优化策略：

预处理父子关系：使用邻接表存储树结构
动态维护存活节点：用vis数组标记节点是否在当前搜索范围内
DFS同时计算权重：在遍历时只计算存活节点的权重

优化后的核心计算逻辑：

def optimized_solution(): import sys sys.setrecursionlimit(10000) n, m = map(int, sys.stdin.readline().split()) w = [0] + list(map(int, sys.stdin.readline().split())) parents = [0, 1] + list(map(int, sys.stdin.readline().split())) # 建树 tree = [[] for _ in range(n+2)] for i in range(2, n+1): tree[parents[i]].append(i) # 处理每个查询 for _ in range(m): target = int(sys.stdin.readline()) vis = [1] * (n+2) # 1表示存活 result = [] root = 1 while True: # DFS计算子树和 W = [0] * (n+2) stack = [(root, False)] while stack: node, processed = stack.pop() if not processed: stack.append((node, True)) for child in reversed(tree[node]): if vis[child]: stack.append((child, False)) else: W[node] = w[node] for child in tree[node]: if vis[child]: W[node] += W[child] # 找最小w_delta min_delta = float('inf') best_node = root stack = [root] while stack: node = stack.pop() delta = abs(2*W[node] - W[root]) if delta < min_delta or (delta == min_delta and node < best_node): min_delta = delta best_node = node for child in tree[node]: if vis[child]: stack.append(child) result.append(str(best_node)) # 检查是否结束 if best_node == target and len(tree[best_node]) == 0: break # 更新存活节点 stack = [best_node] subtree_nodes = set() while stack: node = stack.pop() subtree_nodes.add(node) for child in tree[node]: if vis[child]: stack.append(child) if target in subtree_nodes: # 保留子树 for node in range(1, n+1): if node not in subtree_nodes: vis[node] = 0 root = best_node else: # 移除子树 for node in subtree_nodes: vis[node] = 0 print(' '.join(result))

这个优化版本通过vis数组动态维护存活节点，避免了重复计算。实测下来，对于n=2000的测试用例，运行时间可以控制在500ms以内。

4. 关键算法：树状数组优化

更进一步，我们可以使用树状数组（Fenwick Tree）来优化子树权重的计算。具体思路：

对树进行DFS序编号
使用树状数组维护前缀和
通过区间查询快速计算子树和

实现代码框架：

class FenwickTree: def __init__(self, size): self.n = size self.tree = [0] * (self.n + 2) def update(self, index, delta): while index <= self.n: self.tree[index] += delta index += index & -index def query(self, index): res = 0 while index > 0: res += self.tree[index] index -= index & -index return res def fenwick_solution(): import sys input = sys.stdin.read data = input().split() idx = 0 n = int(data[idx]); idx +=1 m = int(data[idx]); idx +=1 w = [0] + list(map(int, data[idx:idx+n])); idx +=n parents = [0,1] + list(map(int, data[idx:idx+n-1])); idx +=n-1 # 建树并获取DFS序 tree = [[] for _ in range(n+2)] for i in range(2, n+1): tree[parents[i]].append(i) # DFS序编号 in_order = [0]*(n+1) out_order = [0]*(n+1) time = 1 stack = [(1, False)] while stack: node, visited = stack.pop() if not visited: in_order[node] = time time +=1 stack.append((node, True)) for child in reversed(tree[node]): stack.append((child, False)) else: out_order[node] = time -1 # 初始化树状数组 ft = FenwickTree(n) for i in range(1, n+1): ft.update(in_order[i], w[i]) # 处理查询 for _ in range(m): target = int(data[idx]); idx +=1 alive = [True]*(n+1) result = [] while True: # 计算总权重 total = ft.query(n) # 找最佳节点 min_delta = float('inf') best_node = 1 stack = [1] while stack: node = stack.pop() if not alive[node]: continue # 计算子树和 subtree_sum = ft.query(out_order[node]) - ft.query(in_order[node]-1) delta = abs(2*subtree_sum - total) if delta < min_delta or (delta == min_delta and node < best_node): min_delta = delta best_node = node for child in tree[node]: if alive[child]: stack.append(child) result.append(str(best_node)) # 检查是否结束 if best_node == target and not any(alive[child] for child in tree[best_node]): break # 更新存活状态 stack = [best_node] is_in_subtree = False while stack: node = stack.pop() if node == target: is_in_subtree = True if not alive[node]: continue alive[node] = False ft.update(in_order[node], -w[node]) for child in tree[node]: stack.append(child) if is_in_subtree: # 需要保留子树，恢复状态 stack = [best_node] while stack: node = stack.pop() alive[node] = True ft.update(in_order[node], w[node]) for child in tree[node]: stack.append(child) print(' '.join(result))

这种方法的理论时间复杂度是O(nm log n)，在实际测试中表现优异。不过实现起来较为复杂，在考试环境下需要谨慎考虑时间投入与收益比。

5. 实现细节与踩坑记录

在实际编码过程中，有几个容易出错的地方值得注意：

节点编号处理：题目中节点编号从1开始，而Python列表默认从0开始，需要特别注意转换
子树包含判断：需要正确处理目标节点恰好是当前提问节点的情况
权重差计算：wδ = |sum_subtree - (total_sum - sum_subtree)|，这个公式容易写错
多测试用例处理：每个测试用例需要独立处理，不能互相影响

我在调试时遇到的一个典型错误是忘记重置vis数组，导致不同测试用例之间相互影响。修正后的处理方式是：

for _ in range(m): target = int(input()) vis = [1] * (n+2) # 每个测试用例重新初始化 # ...后续处理...

另一个常见问题是递归深度过大导致栈溢出。对于深度较大的树，最好使用非递归的DFS实现：

def dfs_iterative(root, tree, vis): stack = [(root, False)] total = 0 while stack: node, processed = stack.pop() if not processed: stack.append((node, True)) for child in reversed(tree[node]): if vis[child]: stack.append((child, False)) else: total += weights[node] return total