poj1987 Distance Statistics 樹的分治

Description

Frustrated at the number of distance queries required to find a reasonable route for his cow marathon, FJ decides to ask queries from which he can learn more information. Specifically, he supplies an integer K (1 <= K <= 1,000,000,000) and wants to know how many pairs of farms lie at a distance at most K from each other (distance is measured in terms of the length of road required to travel from one farm to another). Please only count pairs of distinct farms (i.e. do not count pairs such as (farm #5, farm #5) in your answer). 

Input

* Lines 1 ..M+1: Same input format as in "Navigation Nightmare" 

* Line M+2: A single integer, K. 

Output

* Line 1: The number of pairs of farms that are at a distance of at most K from each-other. 

樓教主的男人八題裏的，很出名我就簡單說下題意，給你一片森林，每條邊有距離權值，求兩點之間距離小於K的點對數。

這題明顯樹的點分治。每棵樹（包括子樹）的點對分爲三種：root到其他點，跨過root的兩子樹的點，子樹內部的點

因爲分治的思想，其中子樹內部的點我們是不能重複算的，每次計算減掉即可。

dfs求出每棵樹以及子樹的重心作爲新的root（爲了降低複雜度），dfs求出每個點到root的距離，統計即可，其中代碼多爲找重心以及對多棵樹的處理，敲的時候很懵逼，敲完思路就很清晰了。

#include<cstdio>
#include<vector>
#include<cstring>
#include<iostream>
#include<algorithm>
using namespace std;
const int N=40005;
char ch[5];
struct node
{
    int to,w;
    node(int _to,int _w):to(_to),w(_w) {}
};
vector<int> dep;
vector<node> mp[N];
int son[N],d[N],f[N];
bool vis[N],done[N];
int n,m,root,k,size,ans=0;
void dfs(int x) //第一次處理，爲了尋找root後能知道該樹的size值
{
    vis[x]=1;
    son[x]=1;
    for(int i=0; i<mp[x].size(); i++)
    {
        int t=mp[x][i].to;
        if(!vis[t])
        {
            dfs(t);
            son[x]+=son[t];
        }
    }
}
void getroot(int x,int fa) //找root值 注意其實f[x]沒什麼作用,完全可以用兩個變量替代，懶了就這麼寫了
{
	f[x]=0;
	son[x]=1;
	for(int i=0;i<mp[x].size();i++)
	{
		int t=mp[x][i].to;
		if(t!=fa&&!done[t])
		{
			getroot(t,x);
			son[x]+=son[t];
			f[x]=max(f[x],son[t]);
		}
	}
	f[x]=max(f[x],size-son[x]);
	if(f[x]<f[root]) root=x;
}
void getdep(int x,int fa) //找dep
{
	dep.push_back(d[x]);
	for(int i=0;i<mp[x].size();i++)
	{
		int t=mp[x][i].to;
		if(t!=fa&&!done[t])
		{
			d[t]=d[x]+mp[x][i].w;
			getdep(t,x);
		}
	}
}
int calc(int x,int init) //計算root兩邊和root與子樹點距離小於K的點對數量
{
	d[x]=init;
	dep.clear();
	getdep(x,0);
	sort(dep.begin(),dep.end());
	int res=0;
	for(int l=0,r=dep.size()-1;l<r;)
		if(dep[l]+dep[r]<=k) res+=r-l++;
		else r--;
	return res;
}
void work(int x)
{
	ans+=calc(x,0);
	done[x]=1;
	for(int i=0;i<mp[x].size();i++)
	{
		int t=mp[x][i].to;
		if(!done[t])
		{
			ans-=calc(t,mp[x][i].w);
			f[0]=size=son[t];
			getroot(t,root=0);
			work(root);
		}
	}
}
int main()
{
    int tt=0,T,x,y,w;
    scanf("%d %d",&n,&m);
    for(int i=0; i<m; i++)
    {
        scanf("%d %d %d %s",&x,&y,&w,ch);
        mp[x].push_back(node(y,w));
        mp[y].push_back(node(x,w));
    }
    scanf("%d",&k);
    memset(done,0,sizeof(done));
    memset(vis,0,sizeof(vis));
    for(int i=1; i<=n; i++)
        if(!vis[i]) dfs(i);
        ans=0;
	for(int i=1;i<=n;i++)
	{
		if(!done[i])
		{
			f[0]=size=son[i];
			getroot(i,root=0);
			work(root);
		}
	}
	printf("%d\n",ans);
    return 0;
}