Code : The Omnipotent Monster Killer

O(N^3) ↕

Submission

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

const long long inf = 1e18;

class Tree {
  public:
    int n;
    vector<long long> a;
    vector<vector<int>> adj;
    vector<vector<long long>> dp;

    Tree(int n) {
        this->n = n;
        adj.resize(n);
        a.resize(n);
        dp.resize(n, vector<long long>(n + 1, inf));
        // dp[u][r] is the health lost when the game is performed on the subtree
        // rooted at node u and the root is eliminated at round r.
    }

    void dfs(int src, int par) {
        for (auto child : adj[src]) {
            if (child != par) {
                dfs(child, src);
            }
        }

        // Assume parent is killed at round pr.
        // Child can be killed at round cr (!= pr).
        for (int pr = 1; pr <= n; pr++) {
            long long now = a[src] * pr;
            for (auto child : adj[src]) {
                if (child == par) {
                    continue;
                }

                long long mn = inf;
                for (int cr = 1; cr <= n; cr++) {
                    if (cr != pr) {
                        mn = min(mn, dp[child][cr]);
                    }
                }
                now += mn;
            }
            dp[src][pr] = min(dp[src][pr], now);
        }
    }
};

void solve() {
    int n;
    cin >> n;
    Tree t(n);

    for (int i = 0; i < n; i++) {
        cin >> t.a[i];
    }

    for (int i = 0; i < n - 1; i++) {
        int x, y;
        cin >> x >> y;
        x--;
        y--;
        t.adj[x].push_back(y);
        t.adj[y].push_back(x);
    }

    t.dfs(0, -1);
    long long ans = *min_element(t.dp[0].begin(), t.dp[0].end());
    cout << ans << "\n";
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int t;
    cin >> t;
    for (int i = 0; i < t; i++) {
        solve();
    }
    return 0;
}

O(N^2) ↕

Submission

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

const long long inf = 1e18;

class Tree {
  public:
    int n;
    vector<long long> a;
    vector<vector<int>> adj;
    vector<vector<long long>> dp, prefdp, suffdp;

    Tree(int n) {
        this->n = n;
        adj.resize(n);
        a.resize(n);
        dp.resize(n, vector<long long>(n + 1, inf));
        prefdp.resize(n, vector<long long>(n + 2, inf));
        suffdp.resize(n, vector<long long>(n + 2, inf));
        // dp[u][r] is the health lost when the game is performed on the subtree
        // rooted at node u and the root is eliminated at round r.
    }

    void dfs(int src, int par) {
        for (auto child : adj[src]) {
            if (child != par) {
                dfs(child, src);
            }
        }

        // Assume parent is killed at round pr.
        // Child can be killed at round cr (!= pr).
        for (int pr = 1; pr <= n; pr++) {
            long long now = a[src] * pr;
            for (auto child : adj[src]) {
                if (child == par) {
                    continue;
                }
                now += min(prefdp[child][pr - 1], suffdp[child][pr + 1]);
            }
            dp[src][pr] = min(dp[src][pr], now);
        }

        for (int pr = 1; pr <= n; pr++) {
            prefdp[src][pr] = min(prefdp[src][pr - 1], dp[src][pr]);
        }
        for (int pr = n; pr >= 0; pr--) {
            suffdp[src][pr] = min(suffdp[src][pr + 1], dp[src][pr]);
        }
    }
};

void solve() {
    int n;
    cin >> n;
    Tree t(n);

    for (int i = 0; i < n; i++) {
        cin >> t.a[i];
    }

    for (int i = 0; i < n - 1; i++) {
        int x, y;
        cin >> x >> y;
        x--;
        y--;
        t.adj[x].push_back(y);
        t.adj[y].push_back(x);
    }

    t.dfs(0, -1);
    long long ans = *min_element(t.dp[0].begin(), t.dp[0].end());
    cout << ans << "\n";
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int t;
    cin >> t;
    for (int i = 0; i < t; i++) {
        solve();
    }
    return 0;
}

O(N * Log^2(N)) ↕

Submission

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

const long long inf = 1e18;
// max rounds.
const int mxr = 22;

class Tree {
  public:
    int n;
    vector<long long> a;
    vector<vector<int>> adj;
    vector<vector<long long>> dp;

    Tree(int n) {
        this->n = n;
        adj.resize(n);
        a.resize(n);
        dp.resize(n, vector<long long>(mxr + 1, inf));
        // dp[u][r] is the health lost when the game is performed on the subtree
        // rooted at node u and the root is eliminated at round r.
    }

    void dfs(int src, int par) {
        for (auto child : adj[src]) {
            if (child != par) {
                dfs(child, src);
            }
        }

        // Assume parent is killed at round pr.
        // Child can be killed at round cr (!= pr).
        for (int pr = 1; pr <= mxr; pr++) {
            long long now = a[src] * pr;
            for (auto child : adj[src]) {
                if (child == par) {
                    continue;
                }

                long long mn = inf;
                for (int cr = 1; cr <= mxr; cr++) {
                    if (cr != pr) {
                        mn = min(mn, dp[child][cr]);
                    }
                }
                now += mn;
            }
            dp[src][pr] = min(dp[src][pr], now);
        }
    }
};

void solve() {
    int n;
    cin >> n;
    Tree t(n);

    for (int i = 0; i < n; i++) {
        cin >> t.a[i];
    }

    for (int i = 0; i < n - 1; i++) {
        int x, y;
        cin >> x >> y;
        x--;
        y--;
        t.adj[x].push_back(y);
        t.adj[y].push_back(x);
    }

    t.dfs(0, -1);
    long long ans = *min_element(t.dp[0].begin(), t.dp[0].end());
    cout << ans << "\n";
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int t;
    cin >> t;
    for (int i = 0; i < t; i++) {
        solve();
    }
    return 0;
}

O(N * Log(N)) ↕

Submission

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

const long long inf = 1e18;
// max rounds.
const int mxr = 22;

class Tree {
  public:
    int n;
    vector<long long> a;
    vector<vector<int>> adj;
    vector<vector<long long>> dp, prefdp, suffdp;

    Tree(int n) {
        this->n = n;
        adj.resize(n);
        a.resize(n);
        dp.resize(n, vector<long long>(mxr + 1, inf));
        prefdp.resize(n, vector<long long>(mxr + 2, inf));
        suffdp.resize(n, vector<long long>(mxr + 2, inf));
        // dp[u][r] is the health lost when the game is performed on the subtree
        // rooted at node u and the root is eliminated at round r.
    }

    void dfs(int src, int par) {
        for (auto child : adj[src]) {
            if (child != par) {
                dfs(child, src);
            }
        }

        // Assume parent is killed at round pr.
        // Child can be killed at round cr (!= pr).
        for (int pr = 1; pr <= mxr; pr++) {
            long long now = a[src] * pr;
            for (auto child : adj[src]) {
                if (child == par) {
                    continue;
                }
                now += min(prefdp[child][pr - 1], suffdp[child][pr + 1]);
            }
            dp[src][pr] = min(dp[src][pr], now);
        }

        for (int pr = 1; pr <= mxr; pr++) {
            prefdp[src][pr] = min(prefdp[src][pr - 1], dp[src][pr]);
        }
        for (int pr = mxr; pr >= 0; pr--) {
            suffdp[src][pr] = min(suffdp[src][pr + 1], dp[src][pr]);
        }
    }
};

void solve() {
    int n;
    cin >> n;
    Tree t(n);

    for (int i = 0; i < n; i++) {
        cin >> t.a[i];
    }

    for (int i = 0; i < n - 1; i++) {
        int x, y;
        cin >> x >> y;
        x--;
        y--;
        t.adj[x].push_back(y);
        t.adj[y].push_back(x);
    }

    t.dfs(0, -1);
    long long ans = *min_element(t.dp[0].begin(), t.dp[0].end());
    cout << ans << "\n";
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int t;
    cin >> t;
    for (int i = 0; i < t; i++) {
        solve();
    }
    return 0;
}

Almost O(N) ↕

Submission

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

const long long inf = 1e18;
// max rounds.
const int mxr = 22;

class Tree {
  public:
    int n;
    vector<long long> a;
    vector<vector<int>> adj;
    vector<vector<long long>> dp, prefdp, suffdp;

    Tree(int n) {
        this->n = n;
        adj.resize(n);
        a.resize(n);
        dp.resize(n, vector<long long>(mxr + 1, inf));
        prefdp.resize(n, vector<long long>(mxr + 2, inf));
        suffdp.resize(n, vector<long long>(mxr + 2, inf));
        // dp[u][r] is the health lost when the game is performed on the subtree
        // rooted at node u and the root is eliminated at round r.
    }

    void dfs(int src, int par) {
        for (auto child : adj[src]) {
            if (child != par) {
                dfs(child, src);
            }
        }

        int lim = min(mxr, (int)adj[src].size() + 1);

        // Assume parent is killed at round pr.
        // Child can be killed at round cr (!= pr).
        for (int pr = 1; pr <= lim; pr++) {
            long long now = a[src] * pr;
            for (auto child : adj[src]) {
                if (child == par) {
                    continue;
                }
                now += min(prefdp[child][pr - 1], suffdp[child][pr + 1]);
            }
            dp[src][pr] = min(dp[src][pr], now);
        }

        for (int pr = 1; pr <= mxr; pr++) {
            prefdp[src][pr] = min(prefdp[src][pr - 1], dp[src][pr]);
        }
        for (int pr = mxr; pr >= 0; pr--) {
            suffdp[src][pr] = min(suffdp[src][pr + 1], dp[src][pr]);
        }
    }
};

void solve() {
    int n;
    cin >> n;
    Tree t(n);

    for (int i = 0; i < n; i++) {
        cin >> t.a[i];
    }

    for (int i = 0; i < n - 1; i++) {
        int x, y;
        cin >> x >> y;
        x--;
        y--;
        t.adj[x].push_back(y);
        t.adj[y].push_back(x);
    }

    t.dfs(0, -1);
    long long ans = *min_element(t.dp[0].begin(), t.dp[0].end());
    cout << ans << "\n";
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int t;
    cin >> t;
    for (int i = 0; i < t; i++) {
        solve();
    }
    return 0;
}