最大带权子树(hihocoder变种题)

题目来源: http://hihocoder.com/contest/mstest2015jan2/problem/3

本质上这道题是说，给定一个含有N个节点的树，和K个规定节点， 要求找到恰好包含M各节点的子树T， 使得树T在满足包含前面K个规定节点的前提下，使得子树T的权重最大，如果存在则输出权重，不存在这样的子树则输出-1

很容易想到这道题是最大带权子树的变种(本质上是树形DP加01背包)， 但是题目要求除了顶点1，还要求必须包含其他K个顶点(这K个顶点中也可能有1)， 所以题目变得复杂了。。。

稍微思考，可以知道因为一定要求包含顶点1，所以我们在读取输入后，先以1为root重新构建这棵树(代码中build函数就是干这个的)。 然后对于规定的K个顶点，因为他们到顶点1的路径存在且唯一，所以我们找到这些路径，可以知道这些路径恰好构成以1为根的一颗子树S。如果S的顶点数大于M，则说明所求的T一定不存在，则直接返回-1；否则在S的基础上，还可以再选取M - |S| 个节点，构成一个权重更大的子树。

关键是如何选取这个剩下的几个节点，注意到不论最后我们生成的最终子树T是什么样的， 由于树路径的唯一性可知， T如果满足题目要求， 则S一定是T的子树，既然无论如何都包含S，那么我们把S求出来后看做一个新的节点(如顶点0)， 其余不在S中的树的分枝都以这个新顶点为根(这个根权重为S的总权重)。然后对于这颗新的树我们只需要用最大带权子树的方法求出以顶点0为根的且恰好包含M- |S| + 1个节点的最大子树即可，这就是所求的答案。

#include<iostream>
#include<vector>
using namespace std;

#define ll long long int
inline int Max(int a, int b){
	return (a<b?b:a);
}
const int max_n = 102;

int pa[max_n]; //节点父亲
ll val[max_n]; //节点的值
bool is_s[max_n]; //是否已经遍历
vector<int> son[max_n];  //节点的儿子们
 


const int MAX_K = 5;
int must_node[MAX_K]; //必须经过的节点


int N, K, M;

vector<int>::iterator iter;
//重新构建以a为根的子树
void buildtree(int v, int fa){ 
	pa[v] = fa;
	for(vector<int>::iterator iter = son[v].begin(); iter!= son[v].end();){
		if(*iter == fa)
        	iter = son[v].erase(iter);
    	else{
			buildtree(*iter, v);        		
        	++iter;
    	}
	}
}


//下面的数组和dfs是用来求最大带权子树，用的方法是记忆化搜索(伪DP)
//dp[a][left] 表示以点a为根，且恰好有left个节点(包含点a自己)的最大子树的权值
ll dp[max_n][max_n];	

void dfs(const int v, const int M){	
	for(int i = 0 ; i < son[v].size(); i++){
		dfs(son[v][i], M);
		for(int totalM = M; totalM >= 2; totalM--){
			for(int childM = 1; childM < totalM; childM++)
				dp[v][totalM] = Max(dp[v][totalM], dp[v][totalM-childM] + dp[son[v][i]][childM]);
		}
	}
}

int main(){
	//读入输入
	cin >> N >>  K >> M;
	for(int i = 1 ; i <= N; i ++){
		cin >> val[i];
		is_s[i] = false;
	}
	for(int i = 0 ; i < K; i++)
		cin >> must_node[i];
	//读入树的边，由于题目中并没有说边的输入格式，我们只好先默认以无向边的形式读入，然后再重建这个有向树
	for(int i = 0 ; i < N-1; i++){
		int a, b;
		cin >> a >> b;
		son[a].push_back(b);
		son[b].push_back(a);
	}	
	
	//重新构建这颗树，使其有父亲儿子之分并且让1恰好为它的根
	buildtree(1, -1);
	
	//将那棵子树S找到，凡是在S中的点，它们的is_s都是true
	ll sum = 0;   //S中的点的权重之和
	is_s[1]  = true;
	int has_s = 1; //已经遍历过的点的数目
	sum += val[1];
	for(int i = 0 ; i < K; i++){
		int now = must_node[i];
		while(!is_s[now]){
			is_s[now] = true;
			has_s ++;
			sum += val[now];
			now = pa[now];
		}
	}

	//如果子树S太大，则输出-1
	if(has_s > M){
		cout << -1 << endl;
		return 0;
	}
	
	//重构这颗树，即将顶点0看做S缩成的一个点， 权重为S中点的权重之和
	M = M-has_s+1;
	val[0] = sum;
	pa[0] = -1;
	for(int i = 1; i <= N; i++) {
		if(is_s[i]){
			for(int j = 0; j < son[i].size() ; j++){
				if(!is_s[son[i][j]]){
					son[0].push_back(son[i][j]);
					pa[son[i][j]] = 0;
				}
			}
		}
	}
	
	//最后找到这个新树中以0为根的最大子树(恰好有M- has_s + 1)个点
	for(int i = 0 ; i <= N; i++){
		for(int j = 1 ; j <= M; j++)
			dp[i][j] = val[i];
	}	
	dfs(0, M);
	
	cout << dp[0][M] << endl;
	return 0;
}