hihoCoder #1104 Suzhou Adventure

 1 #include <iostream>
 2 #include <vector>
 3 #include <cstring>
 4 #include <set>
 5 
 6 int score[101];
 7 int N, K, M;
 8 int dp[101][101];
 9 std::vector<int> edges[101];
10 int b[101] = {0};
11 int c[101] = {0};
12 std::set<int> must;
13 
14 void makeTree()
15 {
16     int a,b;
17     for (int i = 0; i < N-1; ++i) {
18         std::cin >> a >> b;
19         edges[a].push_back(b);
20         edges[b].push_back(a);
21     }
22 }
23 
24 void convertToBT(int node, int par)
25 {
26     for (int i = 0; i < edges[node].size(); ++i) {
27         if (edges[node][i] != par) {
28             b[edges[node][i]] = c[node];
29             c[node] = edges[node][i];
30             convertToBT(edges[node][i], node);
31         }
32     }
33 }
34 
35 void preProcess(int node, int par)
36 {
37     if (edges[node].size() == 1 && edges[node][0] == par) return;
38 
39     for (int i = 0; i < edges[node].size(); ++i) {
40         if (edges[node][i] != par) {
41             preProcess(edges[node][i], node);
42             if (must.count(edges[node][i])) {
43                 must.insert(node);
44             }
45         }
46     } 
47 }
48 
49 int dfs(int n, int m)
50 {
51     if (dp[n][m] != -1) return dp[n][m];
52     if (m == 0 ) {
53         if (n == 0 || ( must.count(n) == 0 && dfs(c[n],0)==0 && dfs(b[n], 0) == 0)) return dp[n][m] = 0;
54         else return dp[n][m] = -2;
55     }
56     if (n == 0 && m != 0) {return dp[n][m] = -2;}
57     
58 
59     int maxm = -2;
60     if (must.count(n) == 0) {
61         maxm = dfs(b[n], m);
62     }
63 
64     for (int i = 0; i < m; ++i) {
65         if (dfs(c[n], i) == -2 || dfs(b[n], m-1-i) == -2) continue;
66         else {
67             maxm = std::max(maxm, dp[c[n]][i] + dp[b[n]][m-1-i] + score[n]);
68         }
69     }
70 
71     return dp[n][m] = maxm;    
72 
73 }
74 
75 int main()
76 {
77     std::cin >> N >> K >> M;
78     for (int i = 1; i <= N; ++i) {
79         std::cin >> score[i];
80     }
81     int v;
82     for (int i = 0; i < K; ++i) {
83         std::cin >> v;
84         must.insert(v);
85     }
86     if (K > M) {std::cout << -1 << std::endl; return 0;}
87     memset(dp, -1, sizeof dp);
88     makeTree();
89     convertToBT(1,-1);
90     preProcess(1, -1);
91     int maxm = dfs(1, M);
92     if (maxm < 0) std::cout << -1 << std::endl;
93     else std::cout << maxm << std::endl;
94 }