結果

問題 No.2892 Lime and Karin
ユーザー 👑 seekworser
提出日時 2024-07-05 13:47:00
言語 C++23
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 2,599 ms / 8,000 ms
コード長 3,703 bytes
コンパイル時間 5,625 ms
コンパイル使用メモリ 293,240 KB
実行使用メモリ 33,572 KB
最終ジャッジ日時 2024-07-05 21:11:17
合計ジャッジ時間 41,170 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 52
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include <bits/stdc++.h>
#include <atcoder/convolution.hpp>
using namespace std;
using ll = long long;

vector<ll> multiply(vector<ll> a, vector<ll> b) { return atcoder::convolution_ll(a, b); }

int main() {
    ll n; cin >> n;
    vector g(n, vector<ll>(0));
    for (size_t i = 0; i < n - 1; i++) {
        ll u, v; cin >> u >> v;
        --u; --v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    string si; cin >> si;
    vector<ll> s(n);
    for (size_t i=0; i<n; i++) s[i] = (si[i] == '0' ? -1 : 1);

    ll ans = 0;

    vector<bool> removed(n, false);
    vector<ll> sz(n, 0);
    auto centroid_decomposition = [&] (auto self, ll x) -> void {
        auto calc_total_size = [&] (auto self, ll u, ll par) -> ll {
            ll result = 1;
            for (auto v : g[u]) if (v != par && !removed[v]) result += self(self, v, u);
            return result;
        };
        ll total_size = calc_total_size(calc_total_size, x, -1);
        ll m = (total_size + 1) / 2;
        auto find_centroid = [&] (auto self, ll u, ll par) -> pair<ll,ll> {
            ll c = -1; ll mn = n+1;
            sz[u] = 1;
            for (auto v : g[u]) {
                if (v == par || removed[v]) continue;
                auto [cn, mnn] = self(self, v, u);
                if (mnn < mn) {
                    mn = mnn;
                    c = cn;
                }
                sz[u] += sz[v];
            }
            if (sz[u] >= m && sz[u]< mn) {
                mn = sz[u];
                c = u;
            }
            return {c, mn};
        };
        auto [c, _] = find_centroid(find_centroid, x, -1);

        vector<vector<ll>> f(0);
        vector<ll> getas;

        for (auto v : g[c]) {
            if (removed[v]) continue;
            ll mx = 0;
            auto calc_geta = [&] (auto self, ll u, ll par, ll x) -> ll {
                ll result = x;
                mx = max(mx, x);
                for (auto v : g[u]) {
                    if (v == par || removed[v]) continue;
                    result = min(result, self(self, v, u, x+s[v]));
                }
                return result;
            };
            ll geta = max(0LL, -calc_geta(calc_geta, v, c, s[v]));
            f.push_back(vector<ll>(mx+geta+1, 0));
            getas.push_back(geta);
            auto calc_f = [&] (auto self, ll u, ll par, ll x) -> void {
                f[f.size()-1][x+geta] += 1;
                for (auto v : g[u]) {
                    if (v == par || removed[v]) continue;
                    self(self, v, u, x+s[v]);
                }
            };
            calc_f(calc_f, v, c, s[v]);
        }

        ll gmax = 0;
        for (auto gi : getas) gmax = max(gmax, gi);
        vector<ll> total(gmax+sz[x], 0);
        for (size_t i=0; i<f.size(); i++) {
            for (size_t j=0; j<f[i].size(); j++) {
                total[j-getas[i]+gmax] += f[i][j];
            }
        }
        for (size_t i=gmax+1-s[c]; i<total.size(); i++) ans += total[i];

        total = atcoder::convolution_ll(total, total);
        gmax *= 2;
        for (size_t i=0; i<f.size(); i++) {
            f[i] = atcoder::convolution_ll(f[i], f[i]);
            getas[i] *= 2;
            for (size_t j=0; j<f[i].size(); j++) {
                if (f[i][j] != 0) total[j-getas[i]+gmax] -= f[i][j];
            }
        }
        for (size_t i=gmax+1-s[c]; i<total.size(); i++) ans += total[i] / 2;

        removed[c] = true;
        for (auto v : g[c]) { if (!removed[v]) self(self, v); }
        return;
    };
    centroid_decomposition(centroid_decomposition, 0);
    ll add = 0;
    for (size_t i=0; i<n; i++) add += (si[i] - '0');
    cout << ans + add << endl;
    return 0;
}
0