洛谷 P4178 Tree
阿新 • • 發佈:2021-07-22
首先點分治;
由於求解的是小於等於k的個數,於是開一個樹狀陣列維護字首和即可;
#include<bits/stdc++.h> using namespace std; #define LL long long #define MAXN 20009 #define MS 40009 int n,m; struct node{ int to,val; }; vector<node > vc[MS]; int sz[MS],w[MS]; int rt,tr_size; int del[MS]; int dis[MS]; int dislist[MS] ,cntd; int ext[MAXN]; queue<int > Q; LL ans; int p[MAXN]; // 樹狀陣列 int lowbit(int x){ return x&(-x); } void add(int pos,int val){ for(;pos<=20000;pos+=lowbit(pos)) p[pos] += val; } int get_sum(int pos){ int ans = 0; for(;pos;pos-=lowbit(pos)) ans += p[pos]; return ans; } void get_rt(int u,int f){ sz[u] = 1; w[u] = 0; for(auto &nb:vc[u]){ int v = nb.to; if(v != f && !del[v]){ get_rt(v,u); sz[u] += sz[v]; w[u] = max(w[u],sz[v]); } } w[u] = max(w[u],tr_size-sz[u]); if(w[u] < w[rt]) rt = u; } void get_dis(int u,int f){ if(dis[u] <= m) dislist[++cntd] = dis[u]; for(auto &nb:vc[u]){ int v = nb.to; int val = nb.val; if(v != f && !del[v]){ dis[v] = dis[u] + val; get_dis(v,u); } } } void cal(int u){ ext[0]++; add(1,1); // +1方便樹狀陣列維護 for(auto &nb:vc[u]){ int v = nb.to; int val = nb.val; if(!del[v]){ cntd = 0; dis[v] = val; get_dis(v,u); for(int i=1;i<=cntd;i++){ ans += get_sum(m-dislist[i]+1); // 樹狀陣列求字首和 } for(int i=1;i<=cntd;i++){ ext[dislist[i]]++; add(dislist[i]+1,1); Q.push(dislist[i]); } } } ext[0]--; add(1,-1); while(!Q.empty()){ ext[Q.front()]--; add(Q.front()+1,-1); Q.pop(); } } void divide(int u){ del[u] = 1; cal(u); for(auto &nb:vc[u]){ int v = nb.to; if(!del[v]){ w[rt = 0] = tr_size = sz[v]; get_rt(v,0); get_rt(rt,0); divide(rt); } } } void solve(){ cin >> n; for(int i=2;i<=n;i++){ int u,v,val; cin >> u >> v >> val; vc[u].push_back({v,val}); vc[v].push_back({u,val}); } cin >> m; w[rt = 0] = tr_size = n; get_rt(1,0); get_rt(rt,0); divide(rt); cout << ans << "\n"; } int main(){ ios::sync_with_stdio(false); int ce; // cin >> ce; ce = 1; while(ce--){ solve(); } return 0; }