Codeforce 796 c Bank Hacking

Codeforce 796 c Bank Hacking

题目描述：

题目分析：

1. 除了第一次hack之外，以后所有将要hack的bank都要和已经被hack的bank即已经offline的bank有关联；
2. 整个bank之间的关系可以抽象成一课树；
   
   如果当前hack的是父节点，那么只有儿子节点和孙子节点会有影响，即满足题目中邻居和半邻居的关系，他们相应的力量随之加一，对于孙子节点以下的节点则完全没有影响；
3. 由此可以分析得到：
   （1）ans=max，当且仅当max的个数为1，所有的max-1的节点都与max相邻;
   （2）ans=max+1,当且仅当存在一个点满足包括他自己和她所有相邻的节点中值为max的个数为所有max的个数;
   （3）ans=max+2,others;

Code:

注：以下代码存在问题，在cf上提交的时候在第八组数据会超时，可以参考AC代码：https://blog.csdn.net/h1021456873/article/details/70183187
希望求解如何改进我的代码使之不会TLE;

#include <iostream>
#include <climits>
#include <algorithm>
using namespace std;
int a[300010],n,u[300010],v[300010],maxx=INT_MIN,count1=0,count2=0,ans=INT_MAX;
int main() {
	cin>>n;
	for(int i=1; i<=n; i++) {
		cin>>a[i];
	}
	for(int i=1; i<=n-1; i++) {
		cin>>u[i]>>v[i];
	}
	for(int i=1; i<=n; i++) {
		maxx=max(maxx,a[i]);
	}
	for(int i=1; i<=n; i++) {
		if(a[i]==maxx) {
			count1++;
		}
		if(a[i]==maxx-1) {
			count2++;
		}
	}

	if(count1==1) {
		int cnt=0;
		for(int i=1; i<=n-1; i++) {
			if(a[v[i]]==maxx&&a[u[i]]==maxx-1) {
				cnt++;
			} else if(a[u[i]]==maxx&&a[v[i]]==maxx-1) {
				cnt++;
			}
		}
		if(cnt==count2) {
			ans=maxx;
		}
	}
	if(ans==INT_MAX) {
		for(int i=1; i<=n; i++) {
			int cnt=0;
			if(a[i]==maxx) {
				cnt++;
			}
			for(int j=1; j<=n-1; j++) {
				if(v[j]==i&&a[u[j]]==maxx) {
					cnt++;
				} else if(u[j]==i&&a[v[j]]==maxx) {
					cnt++;
				}
			}
			if(cnt==count1) {
				ans=maxx+1;
				break;
			}
		}
	}
	if(ans==INT_MAX) {
		ans=maxx+2;
	}
	cout<<ans<<endl;
	return 0;
}