結果

問題 No.2892 Lime and Karin
ユーザー 👑 seekworser
提出日時 2025-01-30 04:51:29
言語 C++23
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 243 ms / 8,000 ms
コード長 2,530 bytes
コンパイル時間 4,592 ms
コンパイル使用メモリ 290,836 KB
実行使用メモリ 39,168 KB
最終ジャッジ日時 2025-01-30 04:51:43
合計ジャッジ時間 12,376 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 52
権限があれば一括ダウンロードができます

ソースコード

diff # 

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

int main() {
    int n; cin >> n;
    vector g(n, vector<int>());
    for (size_t i=0; i<n-1; i++) {
        int u,v; cin >> u >> v;
        u--; v--;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    string s; cin >> s;
    vector<int> t(n);
    for (size_t i=0; i<n; i++) { t[i] = (s[i] == '1' ? 1 : -1); }
    atcoder::fenwick_tree<ll> fw(2*n+1);  // -n..n
    ll ans = 0;
    vector<int> sz(n, 1);
    vector<int> heavy_child(n, -1);
    auto calc_sz = [&] (auto self, int u, int p) -> int {
        sz[u] = 1;
        for (auto x : g[u]) {
            if (x == p) continue;
            sz[u] += self(self, x, u);
        }
        int mx = -1;
        for (auto x : g[u]) {
            if (x == p) continue;
            if (mx < sz[x]) {
                mx = sz[x];
                heavy_child[u] = x;
            }
        }
        return sz[u];
    };
    calc_sz(calc_sz, 0, -1);

    auto add = [&] (auto self, int u, int p, int val) -> void {
        fw.add(n+val+t[u], 1);
        for (auto x : g[u]) {
            if (x == p) continue;
            self(self, x, u, val+t[u]);
        }
    };
    auto calc = [&] (auto self, int u, int p, int val, int base, int tu) -> void {
        int q = 2 * base - (val + t[u]);
        if (tu == -1) q += 2;
        ans += fw.sum(min(2*n+1, n+q), 2*n+1);
        for (auto x : g[u]) {
            if (x == p) continue;
            self(self, x, u, val+t[u], base, tu);
        }
    };
    auto reset = [&] (auto self, int u, int p, int val) -> void {
        fw.add(n+val+t[u], -1);
        int add = (s[u] == '1' ? 1 : -1);
        for (auto x : g[u]) {
            if (x == p) continue;
            self(self, x, u, val+t[u]);
        }
    };
    auto dfs = [&] (auto self, int u, int p, bool keep, int val) -> void {
        for (auto x : g[u]) {
            if (x == p || x == heavy_child[u]) continue;
            self(self, x, u, false, val+t[u]);
        }
        if (heavy_child[u] != -1) {
            self(self, heavy_child[u], u, true, val+t[u]);
        }
        for (auto x : g[u]) {
            if (x == p || x == heavy_child[u]) continue;
            calc(calc, x, u, val+t[u], val+t[u], t[u]);
            add(add, x, u, val+t[u]);
        }
        fw.add(n+val+t[u], 1);
        ans += fw.sum(n+val+1, 2*n+1);
        if (!keep) {
            reset(reset, u, p, val);
        }
    };
    dfs(dfs, 0, -1, true, 0);
    cout << ans << endl;
}
0