POJ - 1741 (树分治)

题目链接
题目大意

一棵树，n个节点，给出边与边的连接情况以及权值。问有多少条路径的长度<=m。

样例

5 4
1 2 3
1 3 1
1 4 2
3 5 1
8

思路

选定 一个根，我们会发现所有的路径分两种情况：一是经过这个根，二是不经过这个根。对于第二种情况，我们会发现它是第一种情况的子问题。
我们只要想出第一种情况怎么做，然后递归处理第二种情况就可以了。对于我们选定的这个根root，我们求出所有点到它的距离dist[],然后把这些距离排序，问题就变成了在这个序列里面选两个值a,b，使得a + b <= m，有多少种。(想一想怎么做)，其实我们发现扫一遍就能找到有多少种。
但是这样就算多了很多。因为我们要的只是经过这个根的，还要减去不经过它的。这个只要减去经过它子树的即可。
这样我们就需要枚举每一个根然后去计算。
为了防止超时，我们每次用的根都是重心。
时间复杂度是nlognlogn(没听懂为什么)

代码

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <vector>
using namespace std;

const int maxn = 1e6 + 100;
struct node {
   int v, w;
};

vector<node> Edge[maxn];
int sonNum[maxn], sonMax[maxn], dist[maxn], Node[maxn];
int Min, n, m, allNode, id;
bool vis[maxn];

//vis 标记访问没访问过
//sonNum 子树节点的数量 sonMax 子树中最多节点的节点个数
//dist 到当前根的距离
//Node 用来储存所有距离点的数组
//id重心
void dfs(int root, int pre) { //求重心
   sonNum[root] = 1;
   sonMax[root] = 0;
   for(int i = 0; i < Edge[root].size(); i++) {
      int to = Edge[root][i].v;

      if(to == pre || vis[to]) continue;

      dfs(to, root);
      sonNum[root] += sonNum[to];
      sonMax[root] = max(sonMax[root], sonNum[to]);
   }
   sonMax[root] = max(sonMax[root], allNode - sonNum[root]);
   if(sonMax[id] > sonMax[root]) id = root;
}

int ans, cnt;

void GetDist(int root, int pre, int val) {
   dist[root] = val; //这里很重要 仔细想想
   for(int i = 0; i < Edge[root].size(); i++) {
       int to = Edge[root][i].v, w = Edge[root][i].w;
       if(to == pre || vis[to]) continue;
       GetDist(to, root, w + val);
   }
   Node[cnt++] = dist[root];
}

int cal(int root, int val) {
   cnt = 0;
   GetDist(root, 0, val);
   sort(Node, Node + cnt);

   int ans = 0;
   for(int i = 0, j = cnt - 1; i < j; ) {
       if(Node[i] + Node[j] <= m) {
           ans += j - i;
           i++;
       }
       else j--;
   }
   return ans;
}

void solve(int root) {
   ans += cal(root, 0);
   vis[root] = true;

   for(int i = 0; i < Edge[root].size(); i++) {
      int to = Edge[root][i].v, w = Edge[root][i].w;
      if(vis[to]) continue;
      ans -= cal(to, w);
      sonMax[0] = allNode = sonNum[to];
      dfs(to, id = 0);
      solve(id);
   }
}

int main()
{
    while(~scanf("%d %d", &n, &m) && (n || m)) {
         memset(vis, false, sizeof(vis));
         for(int i = 0; i <= n; i++) Edge[i].clear();
         for(int i = 1; i < n; i++) {
            int u, v, w;
            scanf("%d %d %d", &u, &v, &w);
            Edge[u].push_back((node) {v, w});
            Edge[v].push_back((node) {u, w});
         }


         sonMax[0] = allNode = n;
         dfs(1, id = 0);

         ans = 0;
         solve(id);
         printf("%d\n", ans);
    }
}