动态规划-树上dp-1757. 搜集钻石

2020-04-08 08:05:15

问题描述：

蒜国有 n 座城市，编号从 1 到 n，城市间有 n−1 条道路，且保证任意两座城市之间是连通的。每一座城市有一定数量的钻石。

蒜头君想在蒜国搜集钻石。他从城市 1 出发，每天他可以通过城市之间道路开车到另外的城市。当蒜头第一次到一个城市的时候，他可以搜集完这个城市的所有钻石，如果他后面再来到这个城市，就没有砖石可以收集了。

蒜头君只有 k 天时间，请你帮算蒜头君计算他最多可以搜集多少钻石。

num表示每个城市的钻石数量，mp表示城市所有的道路。

样例

样例 1:

输入: n=3,k=2,num=[3,8,3],mp=[[1,3],[3,2]]
输出: 14
解释:
他能得到所有的钻石

样例 2:

输入: n=6,k=2,num=[5,9,8,4,9,2],mp=[[1,6],[6,2],[2,5],[5,3],[5,4]]
输出: 16
解释:
他可以得到第1、2、6个城市中的钻石

注意事项

• 1 <= n <= 100
• 0 <= k <= 200
• nums[i] <= 1000
• 保证任意两座城市可以直接或间接到达

问题求解：

    int m;
    int[] nums;
    int[][][] dp = new int[210][210][2];
    List<Integer>[] graph;
    
    void dfs(int fa,int u)
    {
    	int l=graph[u].size();
    	for(int i=0;i<=m;i++)
    	{
    		dp[u][i][1]=dp[u][i][0]=nums[u - 1];
    	}
    	for(int i=0;i<l;i++)
    	{
    		int v=graph[u].get(i);
    		if(v==fa) continue;
    		dfs(u,v);
    		for(int j=m;j>=1;j--)
    		{
    			for(int e=0;e<=j-1;e++)
    			{
    				if(j-e>=2)
    				{
    					dp[u][j][1]=Math.max(dp[u][j][1],dp[u][e][1]+dp[v][j-e-2][1]);
    					dp[u][j][0]=Math.max(dp[u][j][0],dp[u][e][0]+dp[v][j-e-2][1]);//前面的不回来
    				}
    				dp[u][j][0]=Math.max(dp[u][j][0],dp[u][e][1]+dp[v][j-e-1][0]);//后面的不回来
    			}
    		}
    	}
    }
    
    public int getCount(int n, int k, int[] num, int[][] mp) {
        m = k;
        nums = num;
        graph = new List[n + 1];
        for (int i = 1; i <= n; i++) graph[i] = new ArrayList<>();
        for (int[] e : mp) {
            graph[e[0]].add(e[1]);
            graph[e[1]].add(e[0]);
        }
        dfs(-1, 1);
        return dp[1][m][0];
    }