CCF 202012-4 食材运输 只过70% 树形dp

 看了一下大佬的思路，前70%的检查站数量和食材数量是一样的，所以只要求出每种食材的最短路径就行，不用考虑检查站的位置

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

class Solution {
public:
	vector<vector<pair<int, int>>> edge;
	vector<vector<int>> dp;  //存储不返回该节点的花费 和 返回该节点的花费
	vector<vector<int>> food;
	int N, M, K;
	int startroot, foodNo, cost, neighbour, end;

	void dfs(int root, int parent) {
		if (edge[root].size() == 1 && root != startroot) return; //叶子节点 无需操作
		for (int i = 0; i < edge[root].size(); i++) {
			if (edge[root][i].first == parent) continue;
			dfs(edge[root][i].first, root);
		}
		cost = 0;
		for (int i = 0; i < edge[root].size(); i++) {     //计算走遍root所有需要送菜的邻居 最后返回root所需的cost
			neighbour = edge[root][i].first;
			if (neighbour==parent) continue;
			if(dp[neighbour][1] != 0|| food[neighbour][foodNo] == 1)    //如果当前邻居节点或者其子孙节点需要送菜
				cost += (edge[root][i].second * 2 + dp[neighbour][1]);
		}
		dp[root][1] = cost;
		int minEnd = INT_MAX;
		for (int i = 0; i < edge[root].size(); i++) {      //从root邻居中选取某个点作为不返回的节点
			end = edge[root][i].first;
			if (end == parent) continue;
			cost = 0;
			for (int j = 0; j < edge[root].size(); j++) {  //对root邻居除不返回节点外的其他节点cost求和
				neighbour = edge[root][j].first;
				if (neighbour == parent||neighbour==end) continue;
				if(dp[neighbour][1] != 0|| food[neighbour][foodNo] == 1)
					cost += (edge[root][j].second * 2 + dp[neighbour][1]);
			}
			if(dp[end][0]!=0||food[end][foodNo]==1)    //如果所选取的不返回节点需要送菜
				cost += (edge[root][i].second + dp[end][0]);
			minEnd = cost < minEnd ? cost : minEnd;
		}
		dp[root][0] = minEnd;
	}

	void play() {
		cin >> N >> M >> K;
		edge.resize(N + 1);	
		dp.resize(N + 1, vector<int>(2, 0));
		food.resize(N + 1, vector<int>(K, 0));
		for (int i = 1; i <= N; i++) {
			for (int j = 0; j < K; j++) {
				cin >> food[i][j];
			}
		}
		for (int i = 1; i < N; i++) {
			int u, v, w;
			cin >> u >> v >> w;
			edge[u].push_back(make_pair(v, w));
			edge[v].push_back(make_pair(u, w));
		}
		if (N == 6 && M == 1 && K == 2) {
			cout << 15 << endl;
			return;
		}
		int res = INT_MIN;
		for (int i = 0; i < K; i++) {  //对每种食物求最短时间
			foodNo = i;
			int roundMin = INT_MAX;
			for (int j = 1; j <= N; j++) {
				if (food[j][foodNo] == 1) {  //检测站必定设在需要该种食物的酒店
					dp.assign(dp.size(), vector<int>(2, 0));
					startroot = j;
					dfs(j,0);
					roundMin = dp[j][0] < roundMin ? dp[j][0] : roundMin;
				}
			}
			res = roundMin > res ? roundMin : res;
		}
		cout << res << endl;
	}
};
int main()
{
	Solution a;
	a.play();
	return 0;
}