试题：树上搜索

问题背景

西西艾弗岛大数据中心为了收集用于模型训练的数据，推出了一项自愿数据贡献的系统。岛上的居民可以登录该系统，回答系统提出的问题，从而为大数据中心提供数据。为了保证数据的质量，系统会评估回答的正确性，如果回答正确，系统会给予一定的奖励。

近期，大数据中心需要收集一批关于名词分类的数据。系统中会预先设置若干个名词类别，这些名词类别存在一定的层次关系。例如，“动物”是“生物”的次级类别，“鱼类”是“动物”的次级类别，“鸟类”是“动物”的次级类别，“鱼类”和“鸟类”是“动物”下的邻居类别。这些名词类别可以被按树形组织起来，即除了根类别外，每个类别都有且仅有一个上级类别。
并且所有的名词都可以被归类到某个类别中，即每个名词都有且仅有一个类别与其对应。一个类别的后代类别的定义是：若该类别没有次级类别，则该类别没有后代类别；否则该类别的后代类别为该类别的所有次级类别，以及其所有次级类别的后代类别。

下图示意性地说明了标有星号的类别的次级类别和后代类别。

次级类别与后代类别

系统向用户提出问题的形式是：某名词是否属于某类别，而用户可以选择“是”或“否”来回答问题。该问题的含义是：某名词是否可以被归类到某类别或其后代类别中。

例如，要确定名词“鳕鱼”的类别，系统会向用户提出“鳕鱼是否属于动物”，当用户选择“是”时，系统会进一步询问“鳕鱼是否属于鱼类”，当用户选择“是”时，即可确定“鳕鱼”可以被归类到“鱼类”这一类别。

此外，如果没有更具体的分类，某一名词也可以被归类到非叶子结点的类别中。例如，要确定“猫”的类别，系统可以向用户提出“猫是否属于动物”，当用户选择“是”时，系统会进一步分别询问“猫”是否属于“鱼类”和“鸟类”，当两个问题收到了否定的答案后，系统会确定“猫”的类别是“动物”。

大数据中心根据此前的经验，已经知道了一个名词属于各个类别的可能性大小。为了用尽量少的问题确定某一名词的类别，大数据中心希望小 C 来设计一个方法，以减少系统向用户提出的问题的数量。

问题描述

小 C 观察了事先收集到的数据，并加以统计，得到了一个名词属于各个类别的可能性大小的信息。具体而言，每个类别都可以赋予一个被称为权重的值，值越大，说明一个名词属于该类别的可能性越大。由于每次向用户的询问可以获得两种回答，小 C 联想到了二分策略。他设计的策略如下：

1. 对于每一个类别，统计它和其全部后代类别的权重之和，同时统计其余全部类别的权重之和，并求二者差值的绝对值，计为 W；
2. 选择 W 最小的类别，如果有多个，则选取编号最小的那一个，向用户询问名词是否属于该类别；
3. 如果用户回答“是”，则仅保留该类别及其后代类别，否则仅保留其余类别；
4. 重复步骤 1，直到只剩下一个类别，此时即可确定名词的类别。

小 C 请你帮忙编写一个程序，来测试这个策略的有效性。你的程序首先读取到所有的类别及其上级次级关系，以及每个类别的权重。你的程序需要测试对于被归类到给定类别的名词，按照上述策略提问，向用户提出的所有问题。

样例输入

5 2
10 50 10 10 20
1 1 3 3
5
3

样例输出

2 5
2 5 3 4

 代码

#include <bits/stdc++.h>
using namespace std;

typedef long long LL;

LL totalWeight = 0;
int n, m;
int root;

typedef struct Node {
	int id;
	LL myWeight;
	LL fixedWeight;
	LL familyWeight;
	vector<int> child;

} MyNode;

MyNode data[2001];

// 层次遍历找所有子节点包括自己 
set<int> findMyChildren(int x) {
	queue<int> tempQueue;
	// 存储搜索过的节点
	set<int> childrenSet;
	tempQueue.push(x);
	while (tempQueue.size() > 0) {
		int node = tempQueue.front();
		childrenSet.insert(node);
		MyNode currentNode = data[node];
		int size = currentNode.child.size();
		for (int i = 0; i < size; i++) {
			if (data[currentNode.child[i]].fixedWeight > -1) {
				tempQueue.push(currentNode.child[i]);
			}
		}
		tempQueue.pop();
	}
	return childrenSet;
}

// 深度优先搜索 
LL dfs(int root) {
	MyNode node = data[root];
	LL temp = node.myWeight;
	int size = node.child.size();

	for (int i = 0; i < size; i++) {
		if (data[node.child[i]].fixedWeight > -1) {
			temp += dfs(node.child[i]);
		}
	}
	data[root].familyWeight = temp;
	return temp;
}

void process(int aim) {
	for (int i = 1; i <= n; i++) {
		// 初始化信息，无实际意义 
		data[i].fixedWeight = data[i].myWeight;
		data[i].familyWeight = 0;
	}
	LL total = totalWeight;
	root = 1;
	while (true) {
		LL minWeight = LLONG_MAX;
		int minNode = 0;
		int count = 0;
		// 深度优先更新加权值，需要维护root节点信息
		dfs(root);
		for (int i = 1; i <= n; i++) {
			if (data[i].fixedWeight < 0) continue;
			count++;
			data[i].fixedWeight = abs(total - 2 * data[i].familyWeight);
			if (minWeight > data[i].fixedWeight) {
				minNode = i;
				minWeight = data[i].fixedWeight;
			}
		}

		if (count == 1) return;
		cout << minNode << " ";

		// 查找节点
		set<int> childrenSet = findMyChildren(minNode);
		bool flag = childrenSet.count(aim);
		// flag=true说明目标节点在当前子树，更新头节点信息 
		if (flag) root = minNode;
		for (int i = 1; i <= n; i++) {
			if (data[i].fixedWeight < 0) continue;
			if (flag && !childrenSet.count(data[i].id)) {
				total -= data[i].myWeight;
				data[i].fixedWeight = -1;
			} else if (!flag && childrenSet.count(data[i].id)) {
				total -= data[i].myWeight;
				data[i].fixedWeight = -1;
			}
		}
	}
}

int main() {
	int testSet[101];
	cin >> n >> m;
	for (int i = 1; i <= n; i++) {
		LL temp;
		cin >> temp;
		data[i].id = i;
		data[i].myWeight = temp;
		data[i].fixedWeight = temp;
		data[i].familyWeight = 0;
		totalWeight += temp;
	}
	for (int i = 2; i <= n; i++) {
		int id;
		cin >> id;
		data[id].child.push_back(i);
	}
	for (int i = 0; i < m; i++) {
		cin >> testSet[i];
		process(testSet[i]);
		cout << endl;
	};

	return 0;
}