hdu 5664 Lady CA and the graph （树分治，树状数组）

http://acm.hdu.edu.cn/showproblem.php?pid=5664

题意：定义folded chain为一条路径u，v， u v的lca ！= u && ！= v， 求第k大folded chain

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

#pragma comment(linker, "/STACK:102400000,102400000")
#define eps 1e-9
#define LL long long
#define ULL unsigned long long
#define pii pair<int,int>
#define MP make_pair
#define N (100000 + 10)
#define M (200000 + 10)
#define inf 0x3f3f3f3f
#define mod 1000000007
#define lson ls, ll, md
#define rson rs, md+1, rr


int read() {
char c;
while((c=getchar()) && !('0' <= c && c <= '9'));
int ret = c -'0';
while((c=getchar()) && '0' <=  c && c <= '9') ret = ret * 10 + c - '0';
return ret;
}

int fst[N], vv[M], nxt[M], cost[M], e;

void init() {
memset(fst, -1, sizeof fst);
e = 0;

}

void add(int u, int v, int w) {
vv[e] = v, cost[e] = w, nxt[e] = fst[u], fst[u] = e++;
}

int root, mx[N], sz[N], tot;
bool vis[N];
vector<int> que[N*20], all[N];
vector<int> lab[N];
int cnt;

void dfs1(int u, int p) {
sz[u] = 1;
for(int i = fst[u]; ~i; i= nxt[i]) {
int v = vv[i];
if(v == p || vis[v]) continue;
dfs1(v, u);
sz[u] += sz[v];
}
}

void dfs2(int u, int p) {
mx[u] = 0;
for(int i = fst[u]; ~i; i = nxt[i] ){
int v = vv[i];
if(v == p || vis[v]) continue;
dfs2(v, u);
mx[u] = max(mx[u], sz[v]);
}
mx[u] = max(tot-sz[u], mx[u]);
if(root == -1 || mx[u] < mx[root]) {
root = u;
}
}

void find_root(int &u) {
dfs1(u, -1);
tot = sz[u];
root = -1;
dfs2(u, -1);
u = root;
}

void dfs(int u, int p, int dis) {
all[root].push_back(dis);
que[cnt].push_back(dis);
for(int i = fst[u]; ~i; i = nxt[i]) {
int v = vv[i];
if(v == p || vis[v]) continue;
dfs(v, u, dis+cost[i]);
}
}

int mid;
LL ans;
int n,m;
LL k;
LL cal(vector<int> &que, int k) {
if(k == 0) sort(que.begin(), que.end());
else {
LL ans = 0;
int j = 0;
for(int i = que.size()-1; i >= 0; --i) {
while(j < que.size() && que[j] + que[i] < mid) ++j;
ans += max(0, i-j);
}
return ans;
}
return 0;
}

void solve(int u) {
find_root(u);
vis[u] = 1;
all[u].clear();
lab[u].clear();
all[u].push_back(0);
for(int i = fst[u]; ~i; i = nxt[i] ){
int v = vv[i];
if(vis[v]) continue;
++cnt;
que[cnt].clear();
lab[u].push_back(cnt);
dfs(v, u, cost[i]);
cal(que[cnt], 0);
}
cal(all[u], 0);

for(int i = fst[u]; ~i; i= nxt[i]) {
int v = vv[i];
if(vis[v]) continue;
solve(v);
}
}

int san[N<<2], scnt;
int sum[N];
void f(int u, int p, int dis) {
san[++scnt] = dis;
san[++scnt] = dis-mid;
for(int i = fst[u]; ~i; i = nxt[i]) {
int v = vv[i];
if(v == p) continue;
f(v, u, dis+cost[i]);
}
}

void add(int x, int v) {
while(x <= scnt) {
sum[x] += v;
x += x & -x;
}
}

int query(int x) {
int s = 0;
while(x) {
s += sum[x];
x -= x & -x;
}
return s;
}

int hx(int x) {
return lower_bound(san+1, san+scnt+1, x) - san;
}

void g(int u, int p, int dis){
int id = hx(dis);
ans -= query(hx(dis-mid));
add(id, 1);
for(int i = fst[u]; ~i; i = nxt[i]) {
int v = vv[i];
if(v == p) continue;
g(v, u, dis+cost[i]);
}
add(id, -1);
}

int gao() {
for(int i = 1; i <= n;++i)
ans += cal(all[i], 1);
for(int i = 1; i <= cnt; ++i)
ans -= cal(que[i], 1);
}

bool check() {
ans = 0;
for(int i = 1; i <= n; ++i) vis[i] = 0;
gao();
scnt = 0;
f(m, -1, 0);
sort(san, san+scnt);
scnt = unique(san, san+scnt) - san;
g(m, -1, 0);
return ans >= k;
}

int main() {
int T;
T = read();
while(T--) {
n = read();
m = read();
k = read();
init();
int Max = 0;
for(int i = 1; i < n; ++i) {
int u, v, w;
u = read(); v = read(); w = read();
Max = max(Max, w);
add(u,v, w);
add(v, u, w);
}

memset(vis, 0, sizeof vis);
cnt = 0;
solve(m);

int l = 0, r = n*Max;;
int ret = 0;
while(l < r) {
mid = l + r >> 1;
if(check()) l = mid+1, ret = mid;
else r = mid;
}
if(ret == 0) puts("NO");
else
printf("%d\n", ret);
}
}