POJ1741 Tree (树上点分治/treap+启发式合并)

Description

Give a tree with n vertices,each edge has a length(positive integer less than 1001).
Define dist(u,v)=The min distance between node u and v.
Give an integer k,for every pair (u,v) of vertices is called valid if and only if dist(u,v) not exceed k.
Write a program that will count how many pairs which are valid for a given tree.

---

Input

The input contains several test cases. The first line of each test case contains two integers n, k. (n<=10000) The following n-1 lines each contains three integers u,v,l, which means there is an edge between node u and v of length l.
The last test case is followed by two zeros.

---

Output

For each test case output the answer on a single line.

---

Sample Input

5 4
1 2 3
1 3 1
1 4 2
3 5 1
0 0

---

Sample Output

8

---

Solution

方法一：点分治

关于点分治，首先是找重心。树的重心就是去掉这个点后，所剩下的树中最大的那棵的节点数最小。

如何找重心呢？树形$DP$。这个很简单，略过不提。

然后我们证明重心的性质，就是假如当前树的节点数为$N$，那么去掉重心后最大子树的节点数不超过$N/2$。

如上面那个很丑的图，我们假设$u$是重心，$v$是$u$的最大子树，$w$是路人甲。我们反证：假设$siz[v]>N/2$，那么我们考虑去掉$v$，而不是去掉$u$，记新的大小为$siz'$，那么明显$v$的原先的子树的$siz'<siz[v]$，$siz'[u]<N-N/2=N/2=>siz'[u]<siz[v]$，于是割掉$v$后的最大子树的节点数比割掉$u$后最大子树的大小$siz[v]$还要小，这与$u$是重心矛盾，假设不成立。证毕。

于是如果我们每次这样找到树的重心后递归处理，最多递归$logN$层后就结束了。那我们就这样做呗。每次找到重心后，处理通过根节点的路径，即$dfs$当前子树，然后割掉重心(这里只需要标记一下)，分而治之，对每棵子树执行同样的操作。每一层都有不超过$N$个点被处理到，如果我们处理每一层时要分别排序，需要的时间不会超过$O(NlogN)$，共有$logN$层，总时间就是$O(Nlog^2N)$。

简要的讲述了一下点分治的思想和思路，回到这题。

题目要求$dep[u]+dep[v]-2*dep[lca]<=K$的方案数，我们对当前根向下搜一遍，记录并储存$dep$，然后根据$dep[u]+dep[v]<=K$，直接排序用左右指针单调找，将这个答案记作$ans_1$，然而这样我们不能保证$u,v$的$lca$为当前根，于是我们对于处于同一个儿子的子树自己的答案$ans_2$要减去，不过这里的$dep$要统一加上树边长度，以去掉非法答案。于是这里的答案就包括了$lca$为根的儿子的和所有其他的非法答案，于是当前的正确答案就是$ans_1-\sum ans_2$。在这里还可以记录一下两个点是否属于同一个儿子的子树，然后直接算答案，不过这样其实更麻烦而没有必要。

最后这题解决了，需要注意的是，点分治与CDQ分治和整体二分很像，每一层都是有级别为$N$的元素，分成各个块进行处理，如果每个块需要清空标记之类的，都必须反向标记。因为每个块的时间只能和块内的元素个数(即子树大小)成正比，不能与$N$成正比，不然时间就会退化到$O(N^2)$。

---

方法二：treap+启发式合并

相比于上面的方法，这个方法也挺好。

直接每个点开个$treap$，储存到整棵树的根的深度$dep$，然后在$dfs$过程中将儿子的$treap$往根合并，每次合并时将小的并向大的，保证每次只搜小的那棵$treap$，同时还要在大的$treap$中二叉查找计算答案，就是根据$dep[u]+dep[v]-2*dep[lca]<=K$的算式，在$treap$中维护一个$size$即可。由于这样做，每次小的并向大的，可以证明每个点出现在不超过$logN$个$treap$中，每个点每次的操作不超过$logN$，于是时间复杂度$O(Nlog^2N)$，跟点分治相同。和点分治相比，这种方法更加简单好懂，但点分治也不得不学。

---

Code

Code of 点分治

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <algorithm>
#include <cmath>
#include <iostream>
#define N 10005

using namespace std;
int n, K, head_p[N], cur, dep[N], f[N], sum, root, ans, siz[N], d[N];
bool vis[N];
struct Adj{int next, len, obj;} Edg[N<<1];

void Insert(int a, int b, int c){
    Edg[++cur].next = head_p[a];
    Edg[cur].obj = b;
    Edg[cur].len = c;
    head_p[a] = cur;
} 

void Getroot(int x, int fa){
    siz[x] = 1;  f[x] = 0;
    for(int i = head_p[x]; ~ i; i = Edg[i].next){
        int v = Edg[i].obj;
        if(v == fa || vis[v])  continue;
        Getroot(v, x);
        siz[x] += siz[v];
        f[x] = max(f[x], siz[v]);
    }
    f[x] = max(f[x], sum - siz[x]);
    if(f[x] < f[root])  root = x;
}

void Getdeep(int x, int fa){
    dep[++dep[0]] = d[x];
    for(int i = head_p[x]; ~ i; i = Edg[i].next){
        int v = Edg[i].obj, l = Edg[i].len;
        if(v == fa || vis[v])  continue;
        d[v] = d[x] + l;
        Getdeep(v, x);
    }
}

int Calc(int x, int v){
    d[x] = v;  dep[0] = 0;
    Getdeep(x, 0);
    sort(dep+1, dep+dep[0]+1);
    int l = 1, r = dep[0], ret = 0;
    while(l < r){
        if(dep[l] + dep[r] <= K)  ret += r - l, l ++;
        else  r --;
    }
    return ret;
}

void Solve(int x){
    ans += Calc(x, 0);
    vis[x] = true;
    for(int i = head_p[x]; ~ i; i = Edg[i].next){
        int v = Edg[i].obj, l = Edg[i].len;
        if(vis[v])  continue;
        ans -= Calc(v, l);
        sum = siz[v];
        root = 0;
        Getroot(v, 0);
        Solve(root);
    }
}

int main(){

    freopen("poj1741.in", "r", stdin);
    freopen("poj1741.out", "w", stdout);

    while(~ scanf("%d%d", &n, &K) && n && K){
        cur = -1;
        memset(head_p, -1, sizeof(head_p));
        memset(vis, false, sizeof(vis));

        int a, b, c;
        for(int i = 1; i < n; i++){
            scanf("%d%d%d", &a, &b, &c);
            Insert(a, b, c);
            Insert(b, a, c);
        }

        f[0] = sum = n;
        root = ans = 0;
        Getroot(1, 0);
        Solve(root);
        printf("%d\n", ans);
    }

    return 0;
} 

---

Code of treap+启发式合并

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <cstdlib>
#include <cmath>
#define N 10005

using namespace std;

int n, head_p[N], cur;
int cnt, ans, K, dep[N];
struct Tadj {int next, obj, len;}  Edg[N<<1];

void Init(){
    cur = -1;
    cnt = ans = 0;
    for(int i = 1; i <= n; i++)  head_p[i] = -1;
}

void Insert(int a, int b, int c){
    cur ++;
    Edg[cur].next = head_p[a];
    Edg[cur].obj = b;
    Edg[cur].len = c;
    head_p[a] = cur;
}

struct Treap{
    int val, fix, size;
    Treap *L, *R;
    inline int Lsize() {return L ? L->size : 0;}
    inline int Rsize() {return R ? R->size : 0;}
    void Up(){size = Lsize() + Rsize() + 1;}
}*Root[N], Node[N*20];

Treap *NewTnode(int val){
    Node[cnt].L = Node[cnt].R = NULL;
    Node[cnt].val = val;
    Node[cnt].fix = rand();
    Node[cnt].size = 1;
    return Node+cnt++;
}

void Treap_L_Rot(Treap *&a){
    Treap *b = a->R;
    a->R = b->L;
    b->L = a;
    a = b;
    a->L->Up();
    a->Up();
}
void Treap_R_Rot(Treap *&a){
    Treap *b = a->L;
    a->L = b->R;
    b->R = a;
    a = b;
    a->R->Up();
    a->Up();
}

void Treap_Insert(Treap *&p, int val){
    if(!p)  p = NewTnode(val);
    else if(p->val >= val){
      p->size ++;
      Treap_Insert(p->L, val);
      if(p->fix > p->L->fix)
        Treap_R_Rot(p);
    }
    else{
      p->size ++;
      Treap_Insert(p->R, val);
      if(p->fix > p->R->fix)
        Treap_L_Rot(p);
    }
}
int Treap_Find(Treap *p, int val){
    if(!p)  return 0;
    if(p->val <= val)  return p->Lsize() + 1 + Treap_Find(p->R, val);
    else  return Treap_Find(p->L, val);
}

void work1(Treap *&root, Treap *&p, int deplca){
    if(!p)  return;
    ans += Treap_Find(root, K+2*deplca-p->val);
    work1(root, p->L, deplca);
    work1(root, p->R, deplca);
}

void work2(Treap *&root, Treap *&p){
    if(!p)  return;
    Treap_Insert(root, p->val);
    work2(root, p->L);
    work2(root, p->R);
}

void update(Treap *&root, Treap *&p, int deplca){
    work1(root, p, deplca);
    work2(root, p);
}

void dfs(int root, int fa){
    Root[root] = NULL;
    Treap_Insert(Root[root], dep[root]);
    for(int i = head_p[root]; ~ i; i = Edg[i].next){
      int v = Edg[i].obj, l = Edg[i].len;
      if(v == fa)  continue;
      dep[v] = dep[root] + l;
      dfs(v, root);
      if(Root[root]->size >= Root[v]->size)
        update(Root[root], Root[v], dep[root]);
      else{
        update(Root[v], Root[root], dep[root]);
        Root[root] = Root[v];
      }
    }
}
int main(){

    freopen("poj1741.in", "r", stdin);
    freopen("poj1741.out", "w", stdout);

    while(~ scanf("%d%d", &n, &K) && n && K){

      Init();
      int a, b, c;
      for(int i = 1; i < n; i++){
        scanf("%d%d%d", &a, &b, &c);
        Insert(a, b, c);
        Insert(b, a, c);
      }

      dfs(1, 0);

      printf("%d\n", ans);
    }
    return 0;
}

---

那一世，我转山转水转佛塔啊，不为修来生，只为途中与你相见。