Code: Coloring a Red Black Tree

#include <bits/stdc++.h>

#define int long long
#define db long double

void solve() {
    int n;
    std::cin >> n;

    std::string s;
    std::cin >> s;
    s = "#" + s;

    std::vector<std::vector<int>> edg(n + 1);
    std::vector<int> tot(n + 1);

    for (int i = 0; i < n - 1; ++i) {
        int u, v;
        std::cin >> u >> v;
        edg[u].push_back(v);
        edg[v].push_back(u);
        tot[u]++;
        tot[v]++;
    }

    std::vector<int> r(n + 1);
    for (int i = 1; i <= n; ++i) {
        if (s[i] == '0') {
            for (auto v : edg[i]) {
                if (s[v] == '1') {
                    r[i]++;
                }
            }
        }
    }

    std::vector<std::vector<db>> dp(n + 1, std::vector<db>(2));
    std::vector<bool> vis(n + 1);

    auto dfs = [&](auto &self, int u, int p) -> void {
        vis[u] = true;
        std::vector<int> son;

        for (auto v : edg[u]) {
            if (v != p && s[v] == '0') {
                son.push_back(v);
                self(self, v, u);
            }
        }

        for (int i = 0; i <= 1; ++i) {
            db s1 = 0;
            std::vector<db> temp;

            for (auto v : son) {
                s1 += dp[v][1];
                temp.push_back(dp[v][0] - dp[v][1]);
            }

            std::sort(temp.begin(), temp.end());

            db ans = 1e18, s2 = 0;

            for (int k = 0; k <= son.size(); ++k) {
                int t = r[u] + i + k;

                if (t > 0) {
                    db use = (db)tot[u] / t + s1 + s2;
                    ans = std::min(ans, use);
                }

                if (k < son.size()) {
                    s2 += temp[k];
                }
            }
            dp[u][i] = ans;
        }
    };

    db ans = 0;
    for (int i = 1; i <= n; ++i) {
        if (s[i] == '0' && !vis[i]) {
            dfs(dfs, i, 0);
            ans += dp[i][0];
        }
    }

    std::cout << std::fixed << std::setprecision(12) << ans << "\n";
}

int32_t main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr), std::cout.tie(nullptr);

    int _ = 1;
    std::cin >> _;

    while (_--) {
        solve();
    }

    return 0;
}

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

using ld = long double;
const ld INF = 1e18;
const int BLACK = 0;
const int RED = 1;

void solve() {
    int n;
    cin >> n;
    string s;
    cin >> s;

    vector<vector<int>> adj(n);
    vector<int> degree(n, 0);
    for (int i = 0; i < n - 1; i++) {
        int u, v;
        cin >> u >> v;
        u--;
        v--;
        adj[u].push_back(v);
        adj[v].push_back(u);
        degree[u]++;
        degree[v]++;
    }

    auto color = [&](int u) -> int { return s[u] - '0'; };

    // For each black node, count how many of its neighbors are originally red.
    vector<int> red_nbrs(n, 0);
    for (int u = 0; u < n; u++) {
        if (color(u) == BLACK) {
            for (int v : adj[u]) {
                if (color(v) == RED)
                    red_nbrs[u]++;
            }
        }
    }

    // dp[u][0] = min expected cost for subtree of u when parent is NOT yet red.
    // dp[u][1] = min expected cost for subtree of u when parent IS already red.
    vector<array<ld, 2>> dp(n);
    vector<bool> visited(n, false);

    auto dfs = [&](auto &self, int u, int parent) -> void {
        visited[u] = true;

        // Collect black children in the component tree.
        vector<int> children;
        for (int v : adj[u]) {
            if (v != parent && color(v) == BLACK) {
                children.push_back(v);
                self(self, v, u);
            }
        }

        // Compute dp[u][parent_is_red] for parent_is_red in {0, 1}.
        // Baseline: all children are "late" (processed after u).
        ld baseline = 0;
        for (int c : children)
            baseline += dp[c][RED];

        // Deltas: cost increase if child v is switched from late to early.
        vector<ld> deltas;
        for (int c : children)
            deltas.push_back(dp[c][BLACK] - dp[c][RED]);
        sort(deltas.begin(), deltas.end());

        int m = children.size();

        for (int parent_is_red = 0; parent_is_red <= 1; parent_is_red++) {
            ld best = INF;
            ld delta_prefix = 0;

            for (int k = 0; k <= m; k++) {
                int t = red_nbrs[u] + parent_is_red + k;
                if (t > 0) {
                    ld cost = (ld)degree[u] / t + baseline + delta_prefix;
                    best = min(best, cost);
                }
                if (k < m)
                    delta_prefix += deltas[k];
            }

            dp[u][parent_is_red] = best;
        }
    };

    ld answer = 0;
    for (int u = 0; u < n; u++) {
        if (color(u) == BLACK && !visited[u]) {
            dfs(dfs, u, -1);
            answer += dp[u][0];
        }
    }

    cout << fixed << setprecision(12) << answer << "\n";
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int t;
    cin >> t;
    while (t--)
        solve();
}