第一次接触树分治,看了论文又照挑战上抄的代码,也就理解到这个层次了。。
以后做题中再慢慢体会学习。
题目链接:
http://poj.org/problem?id=1741
题意:
给定树和树边的权重,求有多少对顶点之间的边的权重之和小于等于K。
分析:
树分治。
直接枚举不可,我们将树划分成若干子树。
那么两个顶点有两种情况:
- u,v属于同一子树的顶点对
- u,v属于不同子树的顶点对
第一种情况,对子树递归即可求得。
第二种情况,从u到v的路径必然经过了顶点s,只要先求出每个顶点到s的距离再做统计即可。(注意在第二种情况中减去第一种重复计算的部分)
当树退化成链的形式时,递归的深度则退化为
O(n)
,所以选择每次都找到树的重心作为分隔顶点。重心就是删掉此结点后得到的最大子树的顶点数最少的顶点,删除重心后得到的所有子树顶点数必然不超过
n/2
。
查找重心的时候,假设根为v,先在v的子树中找到一个顶点,使删除该顶点后的最大子树的顶点数最少,然后考虑删除v的情况,获得最大子树的顶点数。
两者选择最小的一个,此时选中的顶点即为重心。
递归的每一层都做了排序
O(nlogn)
,递归深度
O(logn)
,总体时间复杂度
O(nlog2n)
。
代码:
#include<cstdio>
#include<algorithm>
#include<cstring>
#include<vector>
#define sa(a) scanf("%d", &a)
#define mem(a,b) memset(a, b, sizeof(a))
using namespace std;
const int maxn = 1e4 + 5, oo = 0x3f3f3f3f;
int cnt[maxn], vis[maxn];
int K, ans;
struct EDGE{int to; int length;int next;};
int head[maxn];
EDGE edge[maxn * 2];
typedef pair<int, int>pii;
int tot = 0;
void addedge(int u, int v, int l)
{
edge[tot].to = v;
edge[tot].length = l;
edge[tot].next = head[u];
head[u] = tot++;
edge[tot].to = u;
edge[tot].length = l;
edge[tot].next = head[v];
head[v] = tot++;
}
int countsubtree(int v, int p)
{
int ans = 1;
for(int i = head[v]; i != -1; i = edge[i].next){
int w = edge[i].to;
if(w == p||vis[w]) continue;
ans += countsubtree(w, v);
}
return cnt[v] = ans;
}
pii findc(int v, int p, int t)
{
pii res = pii(oo, 0);
int s = 1, m = 1;
for(int i = head[v]; i != -1; i = edge[i].next){
int w = edge[i].to;
if(w == p || vis[w]) continue;
res = min(res, findc(w, v, t));
m = max(m, cnt[w]);
s += cnt[w];
}
m = max(m, t - s);
res = min(res, pii(m, v));
return res;
}
void findpath(int v, int p, int d, vector<int>&ds)
{
ds.push_back(d);
for(int i = head[v]; i != -1; i = edge[i].next){
int w = edge[i].to;
if(w == p || vis[w]) continue;
findpath(w, v, d +edge[i].length, ds);
}
}
int count_pair(vector<int>&ds)
{
int res = 0;
sort(ds.begin(), ds.end());
int j = ds.size() - 1;
int i = 0;
while(i < j){
while(j > i &&ds[i] + ds[j] > K) j--;
res += j - i;
i++;
}
return res;
/*
int j = ds.size();
for(int i = 0; i < ds.size(); i++){
while(j > 0 && ds[i] + ds[j - 1] > K) j--;
res += j - (j > i?1:0);
}
return res / 2;*/
}
void solve(int v)
{
vector<int>ds;
countsubtree(v, -1);
int s = findc(v, -1, cnt[v]).second;
vis[s] =true;
//(1)
for(int i = head[s]; i != -1; i = edge[i].next){
int w = edge[i].to;
if(vis[w]) continue;
solve(w);
}
//(2)
ds.push_back(0);
for(int i = head[s]; i != -1; i = edge[i].next){
int w = edge[i].to;
if(vis[w]) continue;
vector<int>ts;
findpath(w, s, edge[i].length, ts);
ans -= count_pair(ts);
ds.insert(ds.end(), ts.begin(), ts.end());
}
vis[s] = false;
ans += count_pair(ds);
}
void init()
{
tot = 0;
ans = 0;
mem(head, -1);
mem(vis, 0);
mem(cnt, 0);
}
int main (void)
{
int n;
while(scanf("%d%d", &n, &K)== 2 && n + K){
int u, v, l;
init();
for(int i = 0; i < n - 1; i++){
sa(u),sa(v),sa(l);
addedge(u, v, l);
}
solve(1);
printf("%d\n", ans);
}
return 0;
}