結果

問題 No.2892 Lime and Karin
ユーザー hongrock
提出日時 2024-09-13 23:16:15
言語 C++17
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 193 ms / 8,000 ms
コード長 2,130 bytes
コンパイル時間 2,112 ms
コンパイル使用メモリ 198,480 KB
最終ジャッジ日時 2025-02-24 08:12:49
ジャッジサーバーID
（参考情報） 		judge4 / judge6
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 3
other AC * 52
権限があれば一括ダウンロードができます

ソースコード

diff # 

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

#define pb emplace_back
#define mp make_pair
 
using ll = long long;
using pii = pair<int,int>;
 
constexpr int mod = 998244353;
constexpr int inf = 0x3f3f3f3f;
constexpr int N = 1e5 + 10;
constexpr int M = 4e6 + 10;

int n, val[N], tree[N], node_cnt, lch[M], rch[M], sum[M], buf_len, buf[M];
char s[N];
vector<int> G[N];
vector<int> V[N];
ll ans;

int newnode(){
  int o;
  if(buf_len){
    o = buf[--buf_len];
  } else {
    o = ++node_cnt;
  }
  lch[o] = rch[o] = 0;
  sum[o] = 0;
  return o;
}

void add(int &o, int L, int R, int p, int v){
  if(!o)  o = newnode();
  sum[o] += v;
  if(L == R){
    return;
  }
  int mid = (L + R) >> 1;
  if(p <= mid)  add(lch[o], L, mid, p, v);
  else  add(rch[o], mid + 1, R, p, v);
}

int query(int o, int L, int R, int l, int r){
  if(!o)  return 0;
  if(L == l && R == r)  return sum[o];
  int mid = (L + R) >> 1;
  if(r <= mid)  return query(lch[o], L, mid, l, r);
  if(l > mid) return query(rch[o], mid + 1, R, l, r);
  return query(lch[o], L, mid, l, mid) + query(rch[o], mid + 1, R, mid + 1, r);
}

void walk(int o){
  if(!o)  return;
  walk(lch[o]);
  walk(rch[o]);
  buf[buf_len++] = o;
}

void join(int x, int y, int fa){
  if(G[x].size() < G[y].size()){
    swap(G[x], G[y]);
    swap(tree[x], tree[y]);
  }
  for(int v : G[y]){
    int target = val[x] + val[fa] - v + 1;
    if(target <= N){
      ans += query(tree[x], -N, N, max(target, -N), N);
    }
    G[x].pb(v);
  }
  for(int v : G[y]){
    add(tree[x], -N, N, v, 1);
  }
  G[y].clear();
  walk(tree[y]);
}

void dfs(int x, int fa){
  if(s[x-1] == '1'){
    val[x] = 1;
    ++ans;
  } else {
    val[x] = -1;
  }
  val[x] += val[fa];
  tree[x] = newnode();
  add(tree[x], -N, N, val[x], 1);
  G[x].pb(val[x]);
  for(int y : V[x]){
    if(y == fa) continue;
    dfs(y, x);
    join(x, y, fa);
  }
}

void _main(){
  int x, y;
  cin >> n;
  for(int i=1; i<n; ++i){
    cin >> x >> y;
    V[x].pb(y);
    V[y].pb(x);
  }
  cin >> s;
  val[0] = 0;
  dfs(1, 0);
  cout << ans << '\n';
}

int main(){
  ios::sync_with_stdio(0);
  cin.tie(0); cout.tie(0);
  _main();
  return 0;
}
0