[Luogu P4178] Tree

[Luogu P4178] Tree

思路

简单的点分治,用小学二年级奥数难度的容斥原理计算一下,不然会TLE

有点类似于[Luogu P2634] [国家集训队]聪聪可可

代码

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

struct Edge;
struct Node{
	Edge *head;
	int siz, mx, dep;
	bool vis;
}node[40010], *root;
struct Edge{
	Node *to;
	Edge *ne;
	int val;
	Edge(int x, int y, int z) : to(&node[y]), ne(node[x].head), val(z){}
};
int N, K, SumSiz, ans, tmp[40010], cnt;

void getroot(Node *x, Node *fa){
	x->siz=1, x->mx=0;
	for(Edge *i=x->head; i; i=i->ne){
		if(i->to==fa || i->to->vis)	continue;
		getroot(i->to, x);
		x->siz+=i->to->siz;
		x->mx=max(i->to->siz, x->mx);
	}
	x->mx=max(x->mx, SumSiz-x->siz);
	root= root ? (x->mx<root->mx ? x : root) : x;
}

void dfs(Node *x, Node *fa){
	if(x->dep>K)	return;
	tmp[++cnt]=x->dep;
	for(Edge *i=x->head; i; i=i->ne){
		if(i->to==fa || i->to->vis)	continue;
		i->to->dep=x->dep+i->val;
		dfs(i->to, x);
	}
}

int solve(Node *x, int y){
	int sum=0, l=1, r;
	cnt=0;
	x->dep=y;
	dfs(x, NULL);
	r=cnt;
	sort(tmp+1, tmp+r+1);
	while(l<r)
		if(tmp[l]+tmp[r]<=K)
			sum+=r-l, l++;
		else
			r--;
	return sum;
}

void divide(Node *x){
	x->vis=true;
	ans+=solve(x, 0);
	for(Edge *i=x->head; i; i=i->ne){
		if(i->to->vis)	continue;
		ans-=solve(i->to, i->val);
		SumSiz=i->to->siz, root=NULL;
		getroot(i->to, NULL);
		divide(i->to);
	}
}

int main(){
	int x, y, z;
	scanf("%d", &N);
	for(int i=1; i<N; i++){
		scanf("%d%d%d", &x, &y, &z);
		node[x].head=new Edge(x, y, z);
		node[y].head=new Edge(y, x, z);
	}
	scanf("%d", &K);
	SumSiz=N, root=NULL;
	getroot(&node[1], NULL);
	divide(root);
	printf("%d", ans);
	return 0;
}

发表评论

电子邮件地址不会被公开。 必填项已用*标注