【洛谷P5952】水箱

题目

题目链接：https://www.luogu.com.cn/problem/P5952
在地面上有一个水箱，它的俯视图被划分成了 \(n\) 行 \(m\) 列个方格，相邻两个方格之间有一堵厚度可以忽略不计的墙，水箱与外界之间有一堵高度无穷大的墙，因此水不可能漏到外面。已知水箱内每个格子的高度只能是 \([0,H]\) 之间的整数，请统计有多少可能的水位情况。

因为答案可能很大，请对 \(10^9+7\) 取模输出。

我们说两种情况是不同的当且仅当存在至少一个方格的水位在两个情况中不同。

思路

首先考虑两个相邻的水位分别等高的集合，此时如果要合并两个集合，那么显只要水的高度不小于这两个集合之间的边的长度。
那么假设两个集合的水位高度分别为 \(h_1,h_2\)，两个集合内部水位分别不超过 \(h_1,h_2\) 的方案数分别为 \(ans_1,ans_2\)，这两个集合之间的边长度为 \(h\)，那么显然两个集合合并之后方案数为

\[ans=(ans_1+h-h_1)(ans_2+h-h_2) \]

那么将边权从小到大排序后，按照类似最小生成树的方式加边即可。
注意最终要加上水位不超过 \(H\) 的方案数。
时间复杂度 \(O(nm)\)。

代码

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

const int N=500010,MOD=1e9+7;
int n,m,t,tot,father[N],h[N];
ll ans[N];

struct edge
{
	int from,to,dis;
}e[N*2];

int ID(int x,int y)
{
	return (x-1)*m+y;
}

bool cmp(edge x,edge y)
{
	return x.dis<y.dis;
}

int find(int x)
{
	return x==father[x]?x:father[x]=find(father[x]);
}

int main()
{
	scanf("%d%d%d",&n,&m,&t);
	for (int i=1;i<=n;i++)
		for (int j=1,x;j<m;j++)
		{
			scanf("%d",&x);
			e[++tot]=(edge){ID(i,j),ID(i,j+1),x};
		}
	for (int i=1;i<n;i++)
		for (int j=1,x;j<=m;j++)
		{
			scanf("%d",&x);
			e[++tot]=(edge){ID(i,j),ID(i+1,j),x};
		}
	sort(e+1,e+1+tot,cmp);
	for (int i=1;i<=n*m;i++)
		father[i]=i,ans[i]=1LL;
	for (int i=1;i<=tot;i++)
	{
		int x=find(e[i].from),y=find(e[i].to);
		if (x!=y)
		{
			ans[x]=1LL*(ans[x]+e[i].dis-h[x])*(ans[y]+e[i].dis-h[y])%MOD;
			h[x]=e[i].dis; father[y]=x;
		}
	}
	int rt=find(1);
	printf("%lld",(ans[rt]+t-h[rt])%MOD);
	return 0;
}