Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include <bits/stdc++.h>
- using namespace std;
- const int N = 101;
- int K;
- vector<int> tree(N);
- vector<pair<int, int>> v[N];
- int dp[N][N];
- int countTree[N][N];
- int solve(int x, int k) {
- if (dp[x][k] != -1) {
- return dp[x][k];
- }
- // don't make this node sawmill
- int cost1 = 0;
- int tree1 = tree[x];
- for (auto [i, d]: v[x]) {
- cost1 += solve(i, k);
- cost1 += d * countTree[i][k];
- tree1 += countTree[i][k];
- }
- if (k == K || x == 0) {
- countTree[x][k] = tree1;
- cout << x << " " << k << " -> [10] " << cost1 << " " << tree1 << endl;
- return dp[x][k] = cost1;
- }
- // make this node sawmill
- int cost2 = 0;
- int tree2 = 0;
- for (auto [i, d]: v[x]) {
- cost2 += solve(i, k + 1);
- cost2 += d * countTree[i][k + 1];
- tree2 += countTree[i][k + 1];
- }
- if (cost1 < cost2 || (cost1 == cost2 && tree1 <= tree2)) {
- countTree[x][k] = tree1;
- cout << x << " " << k << " -> [1] " << cost1 << " " << tree1 << endl;
- return dp[x][k] = cost1;
- } else {
- countTree[x][k] = tree2;
- cout << x << " " << k << " -> [2] " << cost2 << " " << tree2 << endl;
- return dp[x][k] = cost2;
- }
- }
- int main() {
- int n;
- cin >> n >> K;
- for (int i = 0; i < n; i++) {
- cin >> tree[i];
- }
- for (int i = 0; i < n-1; i++) {
- int x, y, w;
- cin >> x >> y >> w;
- v[y].push_back({x, w});
- }
- for (int i = 0; i < n; i++) {
- for (int j = 0; j <= K; j++) {
- dp[i][j] = -1;
- countTree[i][j] = -1;
- }
- }
- cout << solve(0, 0) << endl;
- }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement