【并查集】luogu_P5952 [POI2018]水箱

题意

一个\(n\times m\)的方格，每个格子的水的高度只能是\([0,H]\)之间的整数，相邻两个方格之间有一堵厚度可以忽略不计的墙，请统计有多少可能的情况(当且仅当存在至少一个方格水的高度在两个情况中不同）。

思路

将墙看成边，把相邻两个格连起来。把一个连通块视为高度相同的块。
现在有两个连通块，它们之间的墙的高度\(h\)（连边）要比这两个联通块之间的最高墙\(h'\)的高度大，不然这两个连通块高度已经是相同的。
所以从小到大枚举墙的高度，类似做最小生成树那样合并连通块，然后统计方案。
\(ans=(ans_u+max_u-h)*(ans_v+max_v-h)\)//每个联通块中加上高度相同的方案，然后是乘法原理

代码

#include <cstdio>
#include <algorithm>
#define P(xx, yy) ((xx - 1) * m + yy)

const int mod = 1e9 + 7;
struct node {
	int u, v, w;
} edge[1000001];
int n, m, h, a, cnt;
int fa[500001], mh[500001], ans[500001];

bool operator < (node x, node y) {
	return x.w < y.w;
}

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

int main() {
	scanf("%d %d %d", &n, &m, &h);
	for (int i = 1; i <= n; i++)
		for (int j = 1; j < m; j++) {
			scanf("%d", &a);
			edge[++cnt] = (node) {P(i, j), P(i, j + 1), a};
		}
	for (int i = 1; i < n; i++)
		for (int j = 1; j <= m; j++) {
			scanf("%d", &a);
			edge[++cnt] = (node) {P(i, j), P(i + 1, j), a};			
		}
	for (int i = 1; i <= n; i++)
		for (int j = 1; j <= m; j++)
			fa[P(i, j)] = P(i, j), ans[P(i, j)] = 1;//0的方案
	std::sort(edge + 1, edge + 1 + cnt);
	for (int i = 1; i <= cnt; i++) {
		int f1 = find(edge[i].u), f2 = find(edge[i].v);
		if (f1 != f2) {
			ans[f2] = (long long)(ans[f1] + edge[i].w - mh[f1]) * (ans[f2] + edge[i].w - mh[f2]) % mod;
			fa[f1] = f2;
			mh[f2] = edge[i].w;
		}
	}
	printf("%d", (ans[find(1)] + h - mh[find(1)]) % mod);
}