Enter The Dragon（并查集）

** Enter The Dragon**
The capital of Ardenia is surrounded by several lakes, and each of them is initially full of water.Currently, heavy rainfalls are expected over the land. Such a rain falls to one of the lakes: if the lake is dry and empty, then it will be filled with water; if the lake is already full, then it will overflow, which will result in a natural disaster. Fortunately, the citizens have a dragon at their disposal (and they will not hesitate to use it). The dragon may drink the whole water from a lake in one sitting. Also, the mages of Ardenia already predicted the weather conditions for the next couple of years. The only question is from which lake and when should the dragon drink to prevent a catastrophe?
Multiple Test Cases
The input contains several test cases. The first line of the input contains a positive integer Z ≤ 40,
denoting the number of test cases. Then Z test cases follow, each conforming to the format described
in section Single Instance Input. For each test case, your program has to write an output conforming to
the format described in section Single Instance Output.
Single Instance Input
The first line of the input instance contains two space-separated positive integers n ≤ 106 and m ≤ 106
,where n is the number of lakes. (There are at most 10 input instances for which n ≥ 105 or m ≥ 105
.) The second line contains the weather forecast for the next m days: m space-separated integers t1, t2, . . . , tm(ti ∈ [0, n]). If ti ∈ [1, n], it means a heavy rainfall over lake ti at day i. If ti = 0, there is no rain at day i, and the dragon has the time to drink the water from one lake of your choice. Note that the dragon does not drink on a rainy day.
Single Instance Output
In the first line your program should output word YES if it is possible to prevent a catastrophic overflow
and NO otherwise. In the former case, you should output the second line containing integers from the range [0, n], where is the number of zeros in the weather forecast description, i.e., the number of
non-rainy days. Each of these integers denotes the number of the lake from which the dragon should
drink; zero means the dragon should not drink from any lake (this might be necessary, as even the dragon cannot drink from an empty lake).
Example
Input Output
4
2 4
0 0 1 1
2 4
0 1 0 2
2 3
0 1 2
2 4
0 0 0 1
NO
YES
1 2
NO
YES
0 1 0

#include<bits/stdc++.h>
using namespace std;

const int MAXN=1e6+7;
int ans[MAXN],fa[MAXN],pre[MAXN],rain[MAXN];
// 快读
inline int read(){
	int x=0,f=1;
	char ch=getchar();
	while(ch<'0' || ch > '9'){ if(ch=='-') f=-1;ch=getchar();
	}
	while(ch <='9' && ch >= '0') x=(x<<3) + (x<<1) + (ch^48) ,ch=getchar();
	return x*f;
}
// 并查集
int find(int x){
	return fa[x]==x?x:fa[x]=find(fa[x]);
}
int m,n,t;
void Dragon(){
//遍历
	for(int i=1;i<=m;i++){
	//下雨天跳过
		if(!rain[i])continue;
	// 寻找i湖前一次下雨时（pre[i]）离得最近的不下雨天(find(pre[i))
		int x=find(pre[rain[i]]);
		// 如果 离得最近的下雨天大于当天，此时肯定是无法将湖清空，直接回复NO
		if(x<i){
		// 然后将本天的不下雨节点忘后推，因为一天只能排空一个湖，如何x+1正好不下雨就直接保存，否则会往后走，找最近的一次
			fa[x] = find(x+1);
		// 保存答案
			ans[x] = rain[i];
		}else {
			printf("NO\n");
			return ;
		}
		// 保存前置节点（即某湖下雨的前一次下雨时间）
		pre[rain[i]] = i;
	}
	printf("YES\n");
	for(int i=1;i<=m;i++)
	if(!rain[i])printf("%d ",ans[i]);
	putchar('\n');
	return ;
}
int main(){
	cin >> t;
	
	while(t--){	
	 memset(pre,0,sizeof(pre));
	 memset(ans,0,sizeof(ans));
	n=read(),m=read();
	for(int i=1;i<=m;i++){
	rain[i] = read();
	}	
	// 将最近的0保存
	int l=m+1;
	fa[m+1]=l;
	for(int i=m;i>=0;i--){
		if(!rain[i] && i!=0) l=i;
		fa[i]=l;
	}
	Dragon();
}
	return 0;
}