CodeChef Easy Draughts

Draughts

去掉无用的叶子节点后再dfs一遍

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<string.h>
using namespace std;
typedef long long LL;
const int N=50005;
int w[N],start[N],visit[N],cnt=1,in[N];
struct edge{
	int v,next;
}e[2*N];
void add(int u,int v){
	e[cnt].v=v;
	e[cnt].next=start[u];
	start[u]=cnt++;
}
LL wtot,ntot,wcnt,ncnt;
void dfs(int x){
	visit[x]=1;
	if(w[x]==1) wcnt++;
	ncnt++;
	for(int i=start[x];i;i=e[i].next){
		int nx=e[i].v;
		if(!visit[nx])
			dfs(nx);
	}
}

int main(){
	int n,m;
	cin>>n>>m;
	for(int i=1;i<=n;i++)
		scanf("%d",&w[i]);
	for(int i=0,u,v;i<m;i++){
		scanf("%d %d",&u,&v);
		add(u,v),add(v,u);
		in[u]++,in[v]++;//入度
	}
	for(int i=1;i<=n;i++){//删除所有不合格节点
		if(in[i]==1&&w[i]==0&&!visit[i]){//叶子节点且不开窗
			int k=i;
			do{
				visit[k]=1;
				for(int j=start[k];j;j=e[j].next){//k的下一点
					if(!visit[e[j].v]){
						k=e[j].v;break;
					}
				}
				in[k]--;
			}while(in[k]==1&&w[k]==0);
		}
	}
	for(int i=1;i<=n;i++){
		wcnt=ncnt=0;
		if(!visit[i]&&in[i]>0){
			dfs(i);
			if(wcnt>1){
				wtot+=wcnt*(wcnt-1)/2;
				ntot+=ncnt;
			}
		}
	}
	cout<<wtot<<" "<<ntot<<endl;
	return 0;
}