[Codeforces 1042E]Vasya and Magic Matrix（期望 DP）

Address

• 洛谷 RemoteJudge
• Codeforces 1042E

Meaning

• 一个 $n$ 行 $m$ 列的矩阵，每个位置有个权值
• 给定一个起始位置 $(r,c)$
• 每次等概率随机地移动到一个权值严格比自己小的位置，得分为移动前和移动后的两个位置的欧几里得距离的平方，这样一直移动直到无法移动为止
• 求得分的期望
• $1\le n,m\le 1000$

Solution

• 非常简单的期望 DP
• 先把权值排序，设排序后第 $i$ 个位置坐标为 $(x_i,y_i)$ ，权值为 $a_i$
• $f[i]$ 表示从第 $i$ 个位置开始移动，直到无法移动的期望得分
• 很容易得出转移
• $f[i]=\frac 1{cnt_i}\sum_{j=1,a_j<a_i}^{i-1}\{f[j]+(x_i-x_j)^2+(y_i-y_j)^2\}$
• 其中 $cnt_i$ 表示权值严格小于 $a_i$ 的位置个数
• 复杂度 $O(n^2m^2)$
• 强行推一波式子
• $f[i]=\frac 1{cnt_i}\sum_{j=1,a_j<a_i}^{i-1}\{f_j+x_j^2+y_j^2-2x_ix_j-2y_iy_j\}+x_i^2+y_i^2$
• $=\frac 1{cnt_i}(\sum_{j=1,a_j<a_i}^{i-1}\{f_j+x_j^2+y_j^2\}-2x_i\sum_{j=1,a_j<a_i}^{i-1}x_j-2y_i\sum_{j=1,a_j<a_i}^{i-1}y_j)+x_i^2+y_i^2$
• 参与转移的 $j$ 是一段前缀且 $j$ 的最大值单调不降
• 可以维护参与转移的 $j$ 的 $f_j+x_j^2+y_j^2$ 、 $x_j$ 、 $y_j$ 之和
• 复杂度 $O(nm)$

Code

#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
#define For(i, a, b) for (i = a; i <= b; i++)

inline int read()
{
	int res = 0; bool bo = 0; char c;
	while (((c = getchar()) < '0' || c > '9') && c != '-');
	if (c == '-') bo = 1; else res = c - 48;
	while ((c = getchar()) >= '0' && c <= '9')
		res = (res << 3) + (res << 1) + (c - 48);
	return bo ? ~res + 1 : res;
}

const int N = 1e6 + 5, ZZQ = 998244353;

int n, m, f[N], inv[N];

struct node
{
	int x, y, val;
} a[N];

inline bool comp(node x, node y)
{
	return x.val < y.val;
}

int main()
{
	int i, j, tn = 0, r, c, p = 1, sum2 = 0, sumx = 0, sumy = 0, sum = 0;
	n = read(); m = read();
	For (i, 1, n) For (j, 1, m)
		c = read(), a[++tn] = {i, j, c};
	n = tn;
	std::sort(a + 1, a + n + 1, comp);
	r = read(); c = read();
	For (i, 1, n) if (a[i].x == r && a[i].y == c) n = i;
	inv[1] = 1;
	For (i, 2, n) inv[i] = 1ll * (ZZQ - ZZQ / i) * inv[ZZQ % i] % ZZQ;
	For (i, 1, n)
	{
		while (p < i && a[p].val < a[i].val)
		{
			sum2 = (1ll * a[p].x * a[p].x +
				1ll * a[p].y * a[p].y + sum2) % ZZQ;
			sumx = (sumx + a[p].x) % ZZQ;
			sumy = (sumy + a[p].y) % ZZQ;
			sum = (sum + f[p]) % ZZQ;
			p++;
		}
		int delta = (1ll * a[i].x * a[i].x +
			1ll * a[i].y * a[i].y) % ZZQ * (p - 1) % ZZQ;
		delta = (delta + sum2) % ZZQ;
		delta = (delta - 2ll * a[i].x * sumx % ZZQ + ZZQ) % ZZQ;
		delta = (delta - 2ll * a[i].y * sumy % ZZQ + ZZQ) % ZZQ;
		f[i] = 1ll * (delta + sum) * inv[p - 1] % ZZQ;
	}
	std::cout << f[n] << std::endl;
	return 0;
}