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);
    }
}