Vasya and Magic Matrix (Codeforces 1042E)
Tags: Math, probabilities

Vasya has got a magic matrix $$$a$$$ of size $$$n \times m$$$. The rows of the matrix are numbered from $$$1$$$ to $$$n$$$ from top to bottom, the columns are numbered from $$$1$$$ to $$$m$$$ from left to right. Let $$$a_{ij}$$$ be the element in the intersection of the $$$i$$$-th row and the $$$j$$$-th column.Vasya has also got a chip. Initially, the chip is in the intersection of the $$$r$$$-th row and the $$$c$$$-th column (that is, in the element $$$a_{rc}$$$). Vasya performs the following process as long as possible: among all elements of the matrix having their value less than the value of the element with the chip in it, Vasya randomly and equiprobably chooses one element and moves his chip to this element.After moving the chip, he adds to his score the square of the Euclidean distance between these elements (that is, between the element in which the chip is now and the element the chip was moved from). The process ends when there are no elements having their values less than the value of the element with the chip in it.Euclidean distance between matrix elements with coordinates $$$(i_1, j_1)$$$ and $$$(i_2, j_2)$$$ is equal to $$$\sqrt{(i_1-i_2)^2 + (j_1-j_2)^2}$$$.Calculate the expected value of the Vasya's final score.It can be shown that the answer can be represented as $$$\frac{P}{Q}$$$, where $$$P$$$ and $$$Q$$$ are coprime integer numbers, and $$$Q \not\equiv 0~(mod ~ 998244353)$$$. Print the value $$$P \cdot Q^{-1}$$$ modulo $$$998244353$$$.

Input: The first line of the input contains two integers $$$n$$$ and $$$m$$$ $$$(1 \le n, m \le 1\,000)$$$ — the number of rows and the number of columns in the matrix $$$a$$$.The following $$$n$$$ lines contain description of the matrix $$$a$$$. The $$$i$$$-th line contains $$$m$$$ integers $$$a_{i1}, a_{i2}, \dots, a_{im} ~ (0 \le a_{ij} \le 10^9)$$$.The following line contains two integers $$$r$$$ and $$$c$$$ $$$(1 \le r \le n, 1 \le c \le m)$$$ — the index of row and the index of column where the chip is now.

Output: Print the expected value of Vasya's final score in the format described in the problem statement.

Note: In the first example, Vasya will move his chip exactly once. The expected value of the final score is equal to $$$\frac{1^2 + 2^2+ 1^2}{3} = 2$$$.

Sample Inputs:
1 4
1 1 2 1
1 3

2 3
1 5 7
2 3 1
1 2

Sample Outputs:
2

665496238

Solution (Python):
import io
import os
import sys
a = os.read(0, os.fstat(0).st_size)
inp = []
i = 0
numb = 0
try:
    while True:
        if a[i] >= b'0'[0]:
            numb = 10 * numb + a[i] - 48
        elif a[i] != b'\r'[0]:
            inp.append(numb)
            numb = 0
        i += 1
except:
    pass
if a and a[-1] >= b'0'[0]:
    inp.append(numb)
ii = 0

def mergesort(A, key=lambda x: x, reverse=False):
    (B, n) = (A[:], len(A))
    for i in range(0, n - 1, 2):
        if key(A[i]) > key(A[i ^ 1]):
            (A[i], A[i ^ 1]) = (A[i ^ 1], A[i])
    width = 2
    while width < n:
        for i in range(0, n, 2 * width):
            (R1, R2) = (min(i + width, n), min(i + 2 * width, n))
            (j, k) = (R1, i)
            while i < R1 and j < R2:
                if key(A[i]) > key(A[j]):
                    B[k] = A[j]
                    j += 1
                else:
                    B[k] = A[i]
                    i += 1
                k += 1
            while i < R1:
                B[k] = A[i]
                k += 1
                i += 1
            while k < R2:
                B[k] = A[k]
                k += 1
        (A, B) = (B, A)
        width *= 2
    if reverse:
        A.reverse()
    return A
MOD = 998244353
MODF = float(MOD)
MAGIC = 6755399441055744.0
SHRT = 65536.0
MODF_INV = 1.0 / MODF
SHRT_INV = 1.0 / SHRT
fround = lambda x: x + MAGIC - MAGIC
fmod = lambda a: a - MODF * fround(MODF_INV * a)
fmul = lambda a, b, c=0.0: fmod(fmod(a * SHRT) * fround(SHRT_INV * b) + a * (b - SHRT * fround(b * SHRT_INV)) + c)

def fpow(x, y):
    if y == 0:
        return 1.0
    res = 1.0
    while y > 1:
        if y & 1 == 1:
            res = fmul(res, x)
        x = fmul(x, x)
        y >>= 1
    return fmul(res, x)
(n, m) = (inp[ii], inp[ii + 1])
ii += 2
a = inp[ii:ii + n * m]
ii += n * m
order = mergesort(list(range(n * m)), key=lambda x: a[x])
(sx, sy) = (inp[ii], inp[ii + 1])
ii += 2
MX = n * m + 1
inv = [0.0] * MX
inv[1] = 1.0
for i in range(2, MX):
    inv[i] = fmul(-1.0 * (MOD // i), inv[MOD % i])
fsum = isum = i2sum = jsum = j2sum = 0
cnt = 0
inv_cnt = 0.0
i = 0
fa = [0] * (n * m)
while i < n * m:
    j = i + 1
    while j < n * m and a[order[i]] == a[order[j]]:
        j += 1
    for k in range(i, j):
        (x, y) = divmod(order[k], m)
        x += 1
        y += 1
        f = fmul(1.0 * cnt, 1.0 * (x * x + y * y), i2sum + j2sum + fsum - 2.0 * (x * isum + y * jsum))
        f = fmul(f, inv_cnt)
        if sx == x and sy == y:
            print(int(f + MODF) % MOD)
            sys.exit()
        fa[k] = int(f)
    for k in range(i, j):
        (x, y) = divmod(order[k], m)
        x += 1.0
        y += 1.0
        fsum = (fa[k] + fsum) % MOD
        isum += x
        i2sum += x * x
        jsum += y
        j2sum += y * y
    cnt += j - i
    inv_cnt = inv[cnt]
    i = j

Solution (Python):
import io
import os
import sys
a = os.read(0, os.fstat(0).st_size)
inp = []
i = 0
numb = 0
try:
    while True:
        if a[i] >= b'0'[0]:
            numb = 10 * numb + a[i] - 48
        elif a[i] != b'\r'[0]:
            inp.append(numb)
            numb = 0
        i += 1
except:
    pass
if a and a[-1] >= b'0'[0]:
    inp.append(numb)
ii = 0

def mergesort(A, key=None, reverse=False):
    if key is None:
        key = lambda x: x
    (B, n) = (A[:], len(A))
    for i in range(0, n - 1, 2):
        if key(A[i]) > key(A[i ^ 1]):
            (A[i], A[i ^ 1]) = (A[i ^ 1], A[i])
    width = 2
    while width < n:
        for i in range(0, n, 2 * width):
            R1 = j = min(i + width, n)
            R2 = min(i + 2 * width, n)
            k = i
            while i < R1 and j < R2:
                if key(A[i]) > key(A[j]):
                    B[k] = A[j]
                    j += 1
                else:
                    B[k] = A[i]
                    i += 1
                k += 1
            while i < R1:
                B[k] = A[i]
                k += 1
                i += 1
            while k < R2:
                B[k] = A[k]
                k += 1
        (A, B) = (B, A)
        width *= 2
    if reverse:
        A.reverse()
    return A
MOD = 998244353
MODF = float(MOD)
MAGIC = 6755399441055744.0
SHRT = 65536.0
MODF_INV = 1.0 / MODF
SHRT_INV = 1.0 / SHRT
fround = lambda x: x + MAGIC - MAGIC
fmod = lambda a: a - MODF * fround(MODF_INV * a)
fmul = lambda a, b, c=0.0: fmod(fmod(a * SHRT) * fround(SHRT_INV * b) + a * (b - SHRT * fround(b * SHRT_INV)) + c)

def fpow(x, y):
    if y == 0:
        return 1.0
    res = 1.0
    while y > 1:
        if y & 1 == 1:
            res = fmul(res, x)
        x = fmul(x, x)
        y >>= 1
    return fmul(res, x)
(n, m) = (inp[ii], inp[ii + 1])
ii += 2
a = inp[ii:ii + n * m]
ii += n * m
order = mergesort(list(range(n * m)), key=lambda x: a[x])
(sx, sy) = (inp[ii], inp[ii + 1])
ii += 2
MX = n * m + 1
inv = [0.0] * MX
inv[1] = 1.0
for i in range(2, MX):
    inv[i] = fmul(-1.0 * (MOD // i), inv[MOD % i])
fsum = isum = i2sum = jsum = j2sum = 0
cnt = 0
inv_cnt = 0.0
i = 0
fa = [0] * (n * m)
while i < n * m:
    j = i + 1
    while j < n * m and a[order[i]] == a[order[j]]:
        j += 1
    for k in range(i, j):
        (x, y) = divmod(order[k], m)
        x += 1
        y += 1
        f = fmul(1.0 * cnt, 1.0 * (x * x + y * y), i2sum + j2sum + fsum - 2.0 * (x * isum + y * jsum))
        f = fmul(f, inv_cnt)
        if sx == x and sy == y:
            print(int(f + MODF) % MOD)
            sys.exit()
        fa[k] = int(f)
    for k in range(i, j):
        (x, y) = divmod(order[k], m)
        x += 1.0
        y += 1.0
        fsum = (fa[k] + fsum) % MOD
        isum += x
        i2sum += x * x
        jsum += y
        j2sum += y * y
    cnt += j - i
    inv_cnt = inv[cnt]
    i = j

Solution (Python):
import io
import os
import sys
a = os.read(0, os.fstat(0).st_size)
inp = []
i = 0
numb = 0
try:
    while True:
        if a[i] >= b'0'[0]:
            numb = 10 * numb + a[i] - 48
        elif a[i] != b'\r'[0]:
            inp.append(numb)
            numb = 0
        i += 1
except:
    pass
if a and a[-1] >= b'0'[0]:
    inp.append(numb)
ii = 0

def mergesort(A, gt_cmp=None):
    if gt_cmp is None:
        gt_cmp = lambda i, j: i > j
    (B, n) = (A[:], len(A))
    for i in range(0, n - 1, 2):
        if gt_cmp(A[i], A[i ^ 1]):
            (A[i], A[i ^ 1]) = (A[i ^ 1], A[i])
    width = 2
    while width < n:
        for i in range(0, n, 2 * width):
            R1 = j = min(i + width, n)
            R2 = min(i + 2 * width, n)
            k = i
            while i < R1 and j < R2:
                if gt_cmp(A[i], A[j]):
                    B[k] = A[j]
                    j += 1
                else:
                    B[k] = A[i]
                    i += 1
                k += 1
            while i < R1:
                B[k] = A[i]
                k += 1
                i += 1
            while k < R2:
                B[k] = A[k]
                k += 1
        (A, B) = (B, A)
        width *= 2
    return A
MOD = 998244353
MODF = float(MOD)
MAGIC = 6755399441055744.0
SHRT = 65536.0
MODF_INV = 1.0 / MODF
SHRT_INV = 1.0 / SHRT
fround = lambda x: x + MAGIC - MAGIC
fmod = lambda a: a - MODF * fround(MODF_INV * a)
fmul = lambda a, b, c=0.0: fmod(fmod(a * SHRT) * fround(SHRT_INV * b) + a * (b - SHRT * fround(b * SHRT_INV)) + c)

def fpow(x, y):
    if y == 0:
        return 1.0
    res = 1.0
    while y > 1:
        if y & 1 == 1:
            res = fmul(res, x)
        x = fmul(x, x)
        y >>= 1
    return fmul(res, x)
(n, m) = (inp[ii], inp[ii + 1])
ii += 2
a = inp[ii:ii + n * m]
ii += n * m
order = mergesort(list(range(n * m)), lambda i, j: a[i] > a[j])
(sx, sy) = (inp[ii], inp[ii + 1])
ii += 2
MX = n * m + 1
inv = [0.0] * MX
inv[1] = 1.0
for i in range(2, MX):
    inv[i] = fmul(-1.0 * (MOD // i), inv[MOD % i])
fsum = isum = i2sum = jsum = j2sum = 0
cnt = 0
inv_cnt = 0.0
i = 0
fa = [0] * (n * m)
while i < n * m:
    j = i + 1
    while j < n * m and a[order[i]] == a[order[j]]:
        j += 1
    for k in range(i, j):
        (x, y) = divmod(order[k], m)
        x += 1
        y += 1
        f = fmul(1.0 * cnt, 1.0 * (x * x + y * y), i2sum + j2sum + fsum - 2.0 * (x * isum + y * jsum))
        f = fmul(f, inv_cnt)
        if sx == x and sy == y:
            print(int(f + MODF) % MOD)
            sys.exit()
        fa[k] = int(f)
    for k in range(i, j):
        (x, y) = divmod(order[k], m)
        x += 1.0
        y += 1.0
        fsum = (fa[k] + fsum) % MOD
        isum += x
        i2sum += x * x
        jsum += y
        j2sum += y * y
    cnt += j - i
    inv_cnt = inv[cnt]
    i = j