結果

問題 No.2892 Lime and Karin
コンテスト
ユーザー pockyny
提出日時 2025-06-30 07:45:43
言語 C++17
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 234 ms / 8,000 ms
コード長 5,588 bytes
コンパイル時間 1,365 ms
コンパイル使用メモリ 91,428 KB
実行使用メモリ 40,576 KB
最終ジャッジ日時 2025-06-30 07:45:54
合計ジャッジ時間 10,943 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 52
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include <iostream>
#include <vector>
using namespace std;
typedef long long ll;
ll INF = 1000000000000000000;
template<typename T = ll>
struct segment_tree{
    int n; vector<T> seg;
    T e(){return 0;}
    T op(T a,T b){return (a + b);}
    segment_tree(){}
    segment_tree(int sz){n = sz; seg.resize(2*n,e());}
    segment_tree(vector<T> v){
        n = v.size(); seg.resize(2*n,e());
        for(int i=0;i<n;i++) seg[i + n] = v[i];
        for(int i=n - 1;i>0;i--) seg[i] = op(seg[i<<1],seg[i<<1|1]);
    }
    void init(int sz){n = sz; seg.resize(2*n);}
    void update(int p,T val){
        for(seg[p += n] += val;p>1;p>>=1) seg[p>>1] = op(seg[p],seg[p^1]);
    }
    T get(int p){return seg[p + n];}
    T query(int l,int r){
        T res = e();
        for(l += n,r += n; l<r;l>>=1,r>>=1){
            if(l&1) res = op(res,seg[l++]);
            if(r&1) res = op(res,seg[--r]);
        }
        return res;
    }
    void clear(){for(int i=0;i<2*n;i++) seg[i] = e();}
};

#include <iostream>
#include <vector>

using namespace std;

// 参考: https://nyaannyaan.github.io/library/tree/dsu-on-tree.hpp.html
template <typename T = int>
struct DSUonTree {
    vector<vector<T>> g;
    vector<int> sz;
    // 部分木クエリにeuler tourで答えるように、iの子要素を列挙する
    vector<int> euler,up,down;
    // eulerとかを途中で使う用のcount
    int idx_;
    // 根
    int root;
    // sub tree計算用
    void dfs1(int cur, int par = -1){
        sz[cur] = 1;
        if ((int)g[cur].size() >= 2 && g[cur][0] == par) swap(g[cur][0], g[cur][1]);
        for(T &v : g[cur]){
            if (v == par) continue;
            dfs1(v,cur);
            sz[cur] += sz[v];
            if (sz[v] > sz[g[cur][0]]) swap(v, g[cur][0]);
        }
    }
    void dfs2(int cur, int par = -1) {
        euler[idx_] = cur;
        down[cur] = idx_++;
        for(auto &v : g[cur]){
            if (v == par) continue;
            dfs2(v, cur);
        }
        up[cur] = idx_;
    }
    
    DSUonTree(vector<vector<T>> &_g,int _root = 0)
      : g(_g),
        sz(_g.size()),
        euler(_g.size()),
        up(_g.size()),
        down(_g.size()),
        idx_(0),
        root(_root){
            dfs1(root);
            dfs2(root);
        }
        
    template <typename UPDATE, typename QUERY, typename CLEAR, typename RESET>
    void run(UPDATE &update, QUERY &query, CLEAR &clear, RESET &reset) {
        auto dfs_update = [&](auto rc,int cur,int par = -1) -> void {
            update(cur);
            for(int u:g[cur]){
                if(u==par) continue;
                rc(rc,u,cur);
            }
            return;
        };
        auto dfs_calc = [&](auto rc,int lca,int cur,int par = -1) -> void {
            query(cur,lca);
            for(int u:g[cur]){
                if(u==par) continue;
                rc(rc,lca,u,cur);
            }
            return;
        };
        auto dsu = [&](auto rc, int cur, int par = -1, bool keep = false) -> void {
            for (int i = 1; i < (int)g[cur].size(); i++){
                if (g[cur][i] != par) rc(rc, g[cur][i], cur, false);
            }
            if(sz[cur]!= 1) rc(rc, g[cur][0], cur, true);
            if (sz[cur] != 1){
                for(int i=1;i<g[cur].size();i++){
                    if(g[cur][i]==par) continue;
                    dfs_calc(dfs_calc,cur,g[cur][i],cur);
                    dfs_update(dfs_update,g[cur][i],cur);
                }
            }
            update(cur);
            query(cur,cur);
            if (!keep){
                for (int i = down[cur]; i < up[cur]; i++) clear(euler[i]);
                reset();
            }
            return;
        };
        dsu(dsu, root);
    }
};


int main(){
    int i,n; cin >> n;
    vector<vector<int>> g(n);
    for(i=0;i<n - 1;i++){
        int x,y; cin >> x >> y; x--; y--;
        g[x].push_back(y); g[y].push_back(x);
    }
    string s; cin >> s;
    vector<int> d(n);
    auto dfs = [&](auto self,int cur,int par,int val) -> void {
        if(s[cur]=='0') val--;
        else val++;
        d[cur] = val;
        for(int u:g[cur]){
            if(u==par) continue;
            self(self,u,cur,val);
        }
    };
    dfs(dfs,0,-1,0);
    // [-n,n]
    segment_tree seg(2*n + 1);
    ll ans = 0;
    // node iの中身をglobalなデータに反映する
    auto update = [&](int i) {
        seg.update(d[i] + n,1);
    };

    // iに対するクエリに答える
    auto query = [&](int i,int lca){
        ll cost = (s[lca]=='1') ? 1 : -1;
        // d[v] > -d[i] + d[lca] - cost なるvの数 
        ll le = max(0LL,-d[i] + 2*d[lca] - cost + n + 1);
        ans += seg.query(le,2*n + 1);
        // cout << "query == " << i + 1 << " " << lca + 1 << " " << seg.query(le,2*n + 1) << "\n";
        // cout << "range == " << -d[i] + d[lca] - cost << "\n";
        // for(int j=0;j<=2*n;j++) cout << seg.get(j) << " ";
        // cout << "\n";
    };

    // seg木とかの場合、一気にresetするのが難しいのでこっちで頂点ごとにclearする
    auto clear = [&](int i) {
        seg.update(d[i] + n,-1);
    };

    // データ構造を空にする
    auto reset = [&]() {};

    DSUonTree<> ds(g);

    ds.run(update,query,clear,reset);
    cout << ans << "\n";
    // for(i=0;i<=2*n;i++) cout << seg.get(i) << " ";
    // cout << "\n";
    // for(i=0;i<n;i++) cout << d[i] << " ";
    // cout << "\n";
    // for(i=0;i<n;i++){
    //     cout << i + 1 << " ";
    //     for(int u:ds.g[i]) cout << u + 1 << " ";
    //     cout << "\n";
    // }
}
0