IOI '05 P6 - Rivers

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

const int INF = 2e9 + 1e8;

int n, k;
vector<vector<int>> children;
vector<int> w, v, d;

// minimum cost of subtree with root, k, with next sawmill being in fixed ancestor
vector<vector<int>> subSubtreeAnswer;

// root, k, minimum cost with root being sawmill
vector<vector<int>> subtreeAnswer;

void subDfs(int node, int ancestor, int distanceToSawmill) {
    subSubtreeAnswer[node] = {w[node]*distanceToSawmill};
    // consider not selecting this node as a sawmill
    for (int child : children[node]) {
        subDfs(child, ancestor, distanceToSawmill + d[child]);
        int a = subSubtreeAnswer[node].size(), b = subSubtreeAnswer[child].size();
        --a, --b;
        int c = min(a+b, k);
        vector<int> newAns(1+c, INF);
        for (int i = 0; i <= a; ++i) {
            for (int j = 0; j <= b && i+j <= c; ++j) {
                newAns[i+j] = min(newAns[i+j], subSubtreeAnswer[node][i] + subSubtreeAnswer[child][j]);
            }
        }
        subSubtreeAnswer[node] = newAns;
    }

    // consider selecting this node as a sawmill
    if (node != ancestor) {
        while (subSubtreeAnswer[node].size() < subtreeAnswer[node].size()) {
            subSubtreeAnswer[node].push_back(INF);
        }
        for (size_t i = 1; i < subtreeAnswer[node].size(); ++i) {
            subSubtreeAnswer[node][i] = min(subSubtreeAnswer[node][i], subtreeAnswer[node][i]);
        }
    }
}

void dfs(int node) {
    for (int child : children[node]) dfs(child);

    subDfs(node, node, 0);

    subtreeAnswer[node].resize(subSubtreeAnswer[node].size()+1);
    for (size_t i = 1; i < subtreeAnswer[node].size(); ++i) subtreeAnswer[node][i] = subSubtreeAnswer[node][i-1];
    // +1 sawmill because this node itself is chosen
}

int main() {
    ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    cin >> n >> k;
    children.resize(1+n);
    w = v = d = vector<int>(1+n);
    for (int i = 1; i <= n; ++i) {
        cin >> w[i] >> v[i] >> d[i];
        children[v[i]].push_back(i);
    }
    ++k; // include choosing the root

    subtreeAnswer.resize(1+n), subSubtreeAnswer.resize(1+n);
    dfs(0);
    cout << subtreeAnswer[0][k] << endl;
    return 0;
}