제출 #1182195

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1182195		Cyanmond	Petrol stations (CEOI24_stations)	C++20	0 / 100	68 ms	13896 KiB

Main

#include <bits/stdc++.h>

using namespace std;
using ll = long long;

void solve() {
    int n;
    ll k;
    cin >> n >> k;
    vector<vector<pair<int, ll>>> g(n);
    for (int i = 0; i < n - 1; ++i) {
        int u, v;
        ll w;
        cin >> u >> v >> w;
        g[u].push_back({v, w});
        g[v].push_back({u, w});
    }

    // centroid decomposition
    vector<bool> seen(n);
    vector<int> sub_size(n);

    auto find_size = [&](auto &&self, int v, int p) -> int {
        sub_size[v] = 1;
        for (auto &[t, w] : g[v]) if (t != p and not seen[t]) {
            auto res = self(self, t, v);
            sub_size[v] += res;
        }
        return sub_size[v];
    };

    auto find_centroid = [&](auto &&self, int v, int p, int s) -> int {
        bool is_centroid = (s - sub_size[v]) <= s / 2;
        for (auto &[t, w] : g[v]) if (t != p and not seen[t]) {
            auto res = self(self, t, v, s);
            if (res != -1) {
                return res;
            }
            if (res > s / 2) {
                res = false;
            }
        }
        return is_centroid ? v : -1;
    };

    vector<ll> ans(n);
    // subtask 5
    static constexpr int D = 10;
    using Data = array<int, D + 1>;

    auto dfs1 = [&](auto &&self, int v, int p, int cof) -> Data {
        Data ret;
        fill(ret.begin(), ret.end(), 0);
        ret[k] = 1;
        for (auto &[t, w] : g[v]) if (t != p and not seen[t]) {
            auto res = self(self, t, v, cof);
            int sum_w = accumulate(res.begin(), res.begin() + w, 0);
            for (int i = 0; i <= k - w; ++i) {
                res[i] = res[i + w];
            }
            fill(res.begin() + (k - w) + 1, res.end(), 0);
            res[k - w] += sum_w;
            for (int i = 0; i <= k; ++i) {
                ret[i] += res[i];
            }
            ans[t] += ll(sum_w) * cof;
        }
        return ret;
    };

    auto dfs2 = [&](auto &&self, int v, int p, Data pbase) -> void {
        ll w_par = 0;
        for (auto &[t, w] : g[v]) if (t == p) {
            w_par = w;
            break;
        }

        int sum_w = accumulate(pbase.begin(), pbase.begin() + w_par, 0);
        for (int i = 0; i <= k - w_par; ++i) {
            pbase[i] = pbase[i + w_par];
        }
        fill(pbase.begin() + (k - w_par) + 1, pbase.end(), 0);
        pbase[k - w_par] += sum_w;
        ans[p] += ll(sum_w) * sub_size[v];

        for (auto &[t, w] : g[v]) if (t != p and not seen[t]) {
            self(self, t, v, pbase);
        }
    };

    auto cent_decomp = [&](auto &&self, int root) -> void {
        // find centroid
        const int s = find_size(find_size, root, -1);
        const int cent = find_centroid(find_centroid, root, -1, s);
        find_size(find_size, cent, -1);

        // do process
        vector<Data> ds(int(g[cent].size()));
        Data r_sum;
        fill(r_sum.begin(), r_sum.end(), 0);
        r_sum[k] = 1;
        for (int i = 0; i < int(g[cent].size()); ++i) {
            auto &[t, w] = g[cent][i];
            if (seen[t]) {
                continue;
            }

            int cof = s - sub_size[t];
            auto res = dfs1(dfs1, t, cent, cof);
            int sum_w = accumulate(res.begin(), res.begin() + w, 0);
            for (int i = 0; i <= k - w; ++i) {
                res[i] = res[i + w];
            }
            fill(res.begin() + (k - w) + 1, res.end(), 0);
            res[k - w] += sum_w;
            ds[i] = res;
            for (int i = 0; i <= k; ++i) {
                r_sum[i] += res[i];
            }

            ans[t] += ll(sum_w) * cof;
        }

        for (int i = 0; i < int(g[cent].size()); ++i) {
            auto &[t, w] = g[cent][i];
            if (seen[t]) {
                continue;
            }
            auto pbase = r_sum;
            for (int j = 0; j <= k; ++j) {
                pbase[j] -= ds[i][j];
            }
            dfs2(dfs2, t, cent, pbase);
        }

        seen[cent] = true;
        for (auto &[t, w] : g[root]) if (not seen[t]) {
            self(self, t);
        }
    };
    cent_decomp(cent_decomp, 0);

    for (auto e : ans) {
        cout << e << '\n';
    }
}

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

    solve();
    return 0;
}

• Subtask 0: 0 / 0
• Subtask 1: 0 / 18
• Subtask 2: 0 / 8
• Subtask 3: 0 / 10
• Subtask 4: 0 / 12
• Subtask 5: 0 / 17
• Subtask 6: 0 / 35