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;
}