生活在树上 小白月赛46

D-生活在树上_牛客小白月赛46 (nowcoder.com)

题意是给你一颗树，然后求每个点走小于等于2的距离最多可以走到多少个点

这题我看着就感觉很不对劲，一开始就无脑暴力，我心里清楚肯定不对，必TLE。

这是TLE的无脑dfs代码：hh

#pragma GCC optimize(1)
#pragma GCC optimize(2)
#pragma GCC optimize(3,"Ofast","inline")
#define IOS ios::sync_with_stdio(false), cin.tie(0);
#include<iostream>
#include<map>
#include<set>
#include<cstdio>
#include<cstring>
#include<vector>
#include<stack>
#include<algorithm>
#include<cmath>
#include<queue>
#include<deque>
using namespace std;
#define int long long 
typedef long long ll;
typedef pair<char,int> PAII;
const int N=2e6+10,M=5050,INF=0x3f3f3f3f,mod=1e9+7;
int h[N],ne[N],e[N],w[N],idx;
bool st[N];
int cnt,sum,pos;
void add(int a,int b,int c)
{
	e[idx]=b;
	w[idx]=c;
	ne[idx]=h[a];
	h[a]=idx++;
}
void dfs(int u)
{
	if(sum>2) return ;
	else 
		if(pos!=u) cnt++; 
	//cout<<u<<" "<<sum<<" "<<cnt<<endl;
	for(int i=h[u];i!=-1;i=ne[i])
	{
		int j=e[i];
		if(!st[j])
		{
	//		cout<<"* "<<j<<endl;
			st[j]=true;
			sum+=w[i];
			dfs(j);
			st[j]=false;
			sum-=w[i];
		}
	}
}
signed main(){
    //IOS; 
    int T;
    T=1;
    //cin>>T;
    while(T--)
    {
    	memset(h,-1,sizeof(h));
		int n;
		cin>>n;
		for(int i=2;i<=n;i++)
		{
			int a,b;
			cin>>a>>b;
			add(i,a,b);
			add(a,i,b);
		}
		for(int i=1;i<=n;i++)
		{
			sum=0,cnt=0;
			for(int i=1;i<=n;i++) st[i]=false;
			pos=i;
			dfs(i);
			cout<<cnt+1<<"\n"; 
		}
    }
    return 0;
}
/*
每次写完判断一下 
*/

但是我好像想不出什么其他方法了，后来看别人代码才知道

就是先可以把每个邻边是1的点的数量找出来，然后遍历所有的点

每个点的边如果边权是1或者2，直接加，但是如果边权是1的话，还可以继续加下一层

就用到了刚刚的所记录的1的数量，然后把到上一层已经算过的-1就好

#pragma GCC optimize(1)
#pragma GCC optimize(2)
#pragma GCC optimize(3,"Ofast","inline")
#define IOS ios::sync_with_stdio(false), cin.tie(0);
#include<iostream>
#include<map>
#include<set>
#include<cstdio>
#include<cstring>
#include<vector>
#include<stack>
#include<algorithm>
#include<cmath>
#include<queue>
#include<deque>
using namespace std;
#define int long long 
typedef long long ll;
typedef pair<char,int> PAII;
const int N=2e6+10,M=5050,INF=0x3f3f3f3f,mod=1e9+7;
int h[N],ne[N],e[N],w[N],idx,cnt[N],res[N];
void add(int a,int b,int c)
{
	e[idx]=b;
	w[idx]=c;
	ne[idx]=h[a];
	h[a]=idx++;
}
signed main(){
    //IOS; 
    int T;
    T=1;
    //cin>>T;
    while(T--)
    {
    	memset(h,-1,sizeof(h));
		int n;
		cin>>n;
		for(int i=2;i<=n;i++)
		{
			int a,b;
			cin>>a>>b;
			add(i,a,b);
			add(a,i,b);
			if(b==1)
			{
				cnt[i]++;
				cnt[a]++;
			}
		}
		for(int i=1;i<=n;i++)
		{
			for(int j=h[i];j!=-1;j=ne[j])
			{
				int k=e[j];
				if(w[j]<=2)
				{
					res[i]++;
					if(w[j]==1) res[i]+=cnt[k]-1;
				}
			}
			cout<<res[i]+1<<"\n";
		}
    }
    return 0;
}
/*
每次写完判断一下 
*/