详解CCF-CSP 202312-3 树上搜索

原题连接

202312-3 树上搜索
在这里插入图片描述

代码及详细注释

//一个树形结构的处理程序，主要用于处理一些权重相关的查询
#include <iostream>
#include <vector>
#include <set>

//定义全局变量
//好处：
//（1）不是局部变量，不用入栈出栈，效率更高
//（2）不用传参，不用拷贝，效率更高
//（3）不需要寻址，效率更高
std::vector<long long> wi, temp; //wi存储当前节点及其子节点的权重，temp存储每个节点的输入权重
std::vector<std::vector<int>> ls; //存储树的邻接列表
std::set<int> yes, no; //yes存储本次查询中可访问且已访问的节点，no存储本次查询中不可访问，即已经被丢弃的节点

//深度优先搜索
void dfs(int node) {
	yes.insert(node); //将当前节点添加到已访问集合中
	for (int i : ls[node]) {
		//如果子节点没有被丢弃
		if (no.find(i) == no.end()) {
			dfs(i); //递归搜索子节点
			wi[node] += wi[i]; //更新当前节点的权重，当前节点的权重等于当前节点的权重加上子节点的权重
		}
	}
}


//获取最小差异权重的节点索引
int getMinIdx(int currentNode){
	int minDiffNode = 1; //初始化最小差异权重的节点索引
	long long minDiff = LLONG_MAX; //初始化最小差异权重
	for (int j : yes) {
		long long diff = std::abs(wi[currentNode] - 2 * wi[j]); //计算以currentNode为根节点的子树的权重和与以j为根节点的子树的权重和的差异权重
		if (diff < minDiff) { //如果差异小于最小差异权重
			minDiff = diff; //更新最小差异权重
			minDiffNode = j; //更新最小差异权重的节点索引
		}
	}
	return minDiffNode; //返回最小差异权重的节点索引
}

//检查目标节点是否在子树中
bool check(int minDiffNode, int targetNode) {
	if(minDiffNode == targetNode) return true; //如果最小差异权重的节点索引等于目标节点，返回true

	for (int i : ls[minDiffNode]) { //遍历最小差异权重的节点的所有子节点
		if (no.find(i) != no.end()) continue; //如果子节点在不可访问集合中，跳过
		if (check(i, targetNode)) return true; //递归检查子节点是否在子树中
	}
	return false; //返回false
}

int main() {

	int n, m; //定义全部的类别和需要查询的类别数量
	std::cin >> n >> m;

	temp.resize(n + 1); //调整临时权重的大小，大小设置为n+1是因为下标从1开始
	ls.resize(n + 1); //调整邻接列表的大小，大小设置为n+1是因为下标从1开始

	//输入每个节点的权重，存储到temp中，下标从1开始
	for (int i = 1; i <= n; i++) {
		std::cin >> temp[i];
	}

	//初始化邻接列表，方便使用[]进行访问
	for (int i = 1; i <= n; i++) {
		ls[i] = std::vector<int>();
	}

	//输入树的邻接列表,下标从2开始,1为根节点,根节点没有父节点
	for (int i = 2; i <= n; i++) {
		int parent;
		std::cin >> parent;
		ls[parent].push_back(i);
	}

	//输入与遍历所有查询节点
	for (int queryIndex = 0; queryIndex < m; ++queryIndex) { // 遍历所有查询
		int targetNode; // 目标节点（类别）
		std::cin >> targetNode; // 输入目标节点
		no.clear(); // 清空不可访问集合

		int currentNode = 1; // 从根节点开始
		while (true) {
			wi = temp; // 重置权重为每个节点的权重
			yes.clear(); // 清空已访问集合
			dfs(currentNode); // 从当前节点开始深度优先搜索，计算当前节点及其子节点的权重和

			// 如果只有一个节点，表明currentNode就是目标节点
			if (yes.size() == 1) {
				break;
			}

			 获取最小差异权重的节点索引
			int minDiffNode = getMinIdx(currentNode); 

			// 如果目标节点在子树中，更新当前节点；否则，将该节点添加到不可访问集合中
			if (check(minDiffNode, targetNode)) {
				currentNode = minDiffNode; // 如果目标节点在子树中，更新当前节点，用于下一次搜索
			}
			else {
				no.insert(minDiffNode); // 否则，将该节点添加到不可访问集合中
			}

			std::cout << minDiffNode << " "; // 输出当前节点索引
		}
		std::cout << "\n"; // 输出换行
	}


	return 0;
}