CF1042E Vasya and Magic Matrix 题解

题意：

瓦西亚得到了一个大小为 $n \times m$ 的魔法矩阵 $a$。矩阵的行从上到下编号为 $1$ 到 $n$，列从左到右编号为 $1$ 到 $m$。设 $a_{ij}$ 是第 $i$ 行和第 $j$ 列的交点处的元素。

瓦西亚还得到了一个芯片。最初，芯片位于第 $r$ 行和第 $c$ 列的交点处（即元素 $a_{rc}$ 处）。瓦西亚执行以下过程，只要可能：在所有矩阵元素中，其值小于芯片所在元素的值，瓦西亚随机且等概率地选择一个元素，并将芯片移动到该元素。

移动芯片后，他将其得分增加到这些元素之间的欧几里德距离的平方（即芯片现在所在的元素与芯片移动前所在的元素之间的距离）。当没有元素的值小于芯片所在元素的值时，过程结束。

矩阵元素之间的欧几里德距离，其坐标分别为 $(i_1, j_1)$ 和 $(i_2, j_2)$，等于 $\sqrt{(i_1-i_2)^2 + (j_1-j_2)^2}$。

计算瓦西亚最终得分的期望值。

可以证明答案可以表示为 $\frac{P}{Q}$，其中 $P$ 和 $Q$ 是互质的整数，并且 $Q \not\equiv 0\pmod{998244353}$。以 $P \cdot Q^{-1} \pmod{998244353}$ 的形式输出值。

输入格式：

输入的第一行包含两个整数 $n$ 和 $m$ $(1 \le n, m \le 1,000)$ —— 矩阵 $a$ 的行数和列数。

接下来的 $n$ 行包含矩阵 $a$ 的描述。第 $i$ 行包含 $m$ 个整数 $a_{i1}, a_{i2}, \dots, a_{im}$ $(0 \le a_{ij} \le 10^9)$。

接下来的一行包含两个整数 $r$ 和 $c$ $(1 \le r \le n, 1 \le c \le m)$ —— 芯片当前所在的行号和列号。

输出格式：

以问题陈述中描述的格式输出瓦西亚最终得分的期望值。

因为一个点 $i$ 只会被权值比 $a_i$ 小的 $j$ 贡献，所以我们可以先把所有的点按权值顺序排序。

设 $f_i$ 表示第 $i$ 个点作为起点的期望，记有 $T$ 个点的权值小于 $a_i$ 则

$f_i = \frac{\sum_{a_j<a_i}\left( f_{j}+(x_i-x_j)^2+(y_i-y_j)^2\right)}{T}$

因为已经排序过了，所以这个式子可以很好的前缀和优化，不过需要先把距离拆开来

$\begin{gathered} \left(x_i-x_j\right)^2+\left(y_i-y_j\right)^2=x_i^2-2 \times x_i \times x_j+x_j^2+y_i^2-2 \times y_i \times y_j+y_j^2 \\ =x_i^2+x_j^2+y_i^2+y_j^2-2 \times\left(x_i \times x_j+y_i \times y_j\right) \end{gathered}$

于是我们只要维护 $\sum x_j,~\sum y_j,~\sum(x_j^2+y_j^2),~\sum f_j$ 这四个量就可以 $\mathcal{O}(\log V)$ 转移了

时间复杂度 $\mathcal{O}(nm \log V)$

代码：

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF 0x3f3f3f3f3f3f3f3f
const int mod = 998244353;
void up(int &x,int y) { x < y ? x = y : 0; }
void down(int &x,int y) { x > y ? x = y : 0; }
void add(int &x,int y) { (x += y) >= mod ? x -= mod : 0; }
#define N ((int)(1e6 + 15))
#define sq(x) ((x) * (x))

int tot, sumx, sumy, sump, sumdp, dp[N];
struct node { int x, y, v; } p[N];
int qpow(int a, int b)
{
    int r = 1;
    while(b)
    {
        if(b & 1) r = r * a % mod;
        b >>= 1; a = a * a % mod;
    }
    return r;
}
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);
    // freopen("check.in","r",stdin);
    // freopen("check.out","w",stdout);
    int n, m; cin >> n >> m;
    for(int i = 1; i <= n; i++)
        for(int j = 1, x; j <= m; j++) { cin >> x, p[++tot] = {i, j, x}; }
    int sx, sy; cin >> sx >> sy;
    sort(p + 1, p + 1 + tot, [](node a, node b) { return a.v < b.v; });
    for(int i = 1, l = 1; i <= tot; i++)
    {
        dp[i] = 0;
        if(p[i].v != p[l].v) {
            while(l < i) {
                add(sumx, p[l].x), add(sumy, p[l].y);
                add(sump, sq(p[l].x) + sq(p[l].y)); add(sumdp, dp[l]); ++l;
            }
        }
        add(dp[i], (l - 1) * (sq(p[i].x) + sq(p[i].y)) % mod);
        add(dp[i], mod - 2 * (p[i].x * sumx % mod + p[i].y * sumy % mod) % mod);
        add(dp[i], (sump + sumdp) % mod);
        dp[i] = dp[i] * qpow(l - 1, mod - 2) % mod;
        if(p[i].x == sx && p[i].y == sy) return cout << dp[i] << '\n', 0;
    }
    return 0;
}