Magic Square (Codeforces 1864G)
Tags: Implement, combinatorics

Aquamoon has a Rubik's Square which can be seen as an $$$n \times n$$$ matrix, the elements of the matrix constitute a permutation of numbers $$$1, \ldots, n^2$$$.Aquamoon can perform two operations on the matrix:  Row shift, i.e. shift an entire row of the matrix several positions (at least $$$1$$$ and at most $$$n-1$$$) to the right. The elements that come out of the right border of the matrix are moved to the beginning of the row. For example, shifting a row $$$\begin{pmatrix} a & b & c \end{pmatrix}$$$ by $$$2$$$ positions would result in $$$\begin{pmatrix} b & c & a \end{pmatrix}$$$; Column shift, i.e. shift an entire column of the matrix several positions (at least $$$1$$$ and at most $$$n-1$$$) downwards. The elements that come out of the lower border of the matrix are moved to the beginning of the column. For example, shifting a column $$$\begin{pmatrix} a \\ b \\ c \end{pmatrix}$$$ by $$$2$$$ positions would result in $$$\begin{pmatrix} b\\c\\a \end{pmatrix}$$$. The rows are numbered from $$$1$$$ to $$$n$$$ from top to bottom, the columns are numbered from $$$1$$$ to $$$n$$$ from left to right. The cell at the intersection of the $$$x$$$-th row and the $$$y$$$-th column is denoted as $$$(x, y)$$$.Aquamoon can perform several (possibly, zero) operations, but she has to obey the following restrictions: each row and each column can be shifted at most once; each integer of the matrix can be moved at most twice; the offsets of any two integers moved twice cannot be the same. Formally, if integers $$$a$$$ and $$$b$$$ have been moved twice, assuming $$$a$$$ has changed its position from $$$(x_1,y_1)$$$ to $$$(x_2,y_2)$$$, and $$$b$$$ has changed its position from $$$(x_3,y_3)$$$ to $$$(x_4,y_4)$$$, then $$$x_2-x_1 \not\equiv x_4-x_3 \pmod{n}$$$ or $$$y_2-y_1 \not\equiv y_4-y_3 \pmod{n}$$$.Aquamoon wonders in how many ways she can transform the Rubik's Square from the given initial state to a given target state. Two ways are considered different if the sequences of applied operations are different. Since the answer can be very large, print the result modulo $$$998\,244\,353$$$.

Input: Each test contains multiple test cases. The first line contains the number of test cases $$$t$$$ ($$$1 \le t \le 2\cdot 10^4$$$). The description of the test cases follows.The first line of each test case contains an integer $$$n$$$ ($$$3\le n \le 500$$$).The $$$i$$$-th of the following $$$n$$$ lines contains $$$n$$$ integers $$$a_{i1}, \ldots, a_{in}$$$, representing the $$$i$$$-th row of the initial matrix ($$$1 \le a_{ij} \le n^2$$$).The $$$i$$$-th of the following $$$n$$$ lines contains $$$n$$$ integers $$$b_{i1}, \ldots, b_{in}$$$, representing the $$$i$$$-th row of the target matrix ($$$1 \le b_{ij} \le n^2$$$).It is guaranteed that both the elements of the initial matrix and the elements of the target matrix constitute a permutation of numbers $$$1, \ldots, n^2$$$.It is guaranteed that the sum of $$$n^2$$$ over all test cases does not exceed $$$250\,000$$$.

Output: For each test case, if it is possible to convert the initial state to the target state respecting all the restrictions, output one integer — the number of ways to do so, modulo $$$998\,244\,353$$$.If there is no solution, print a single integer $$$0$$$.

Note: In the first test case, the only way to transform the initial matrix to the target one is to shift the second row by $$$1$$$ position to the right, and then shift the first column by $$$1$$$ position downwards.In the second test case, it can be shown that there is no correct way to transform the matrix, thus, the answer is $$$0$$$.

Sample Inputs:
4
3
1 2 3
4 5 6
7 8 9
7 2 3
1 4 5
6 8 9
3
1 2 3
4 5 6
7 8 9
3 2 1
6 5 4
9 7 8
3
1 2 3
4 5 6
7 8 9
7 8 1
2 3 4
5 6 9
3
1 2 3
4 5 6
7 8 9
3 8 4
5 1 9
7 6 2

Sample Outputs:
1
0
0
4

Solution (Python):
mod = 998244353

def fac(n, memo={}):
    if n in memo:
        return memo[n]
    if n == 0:
        return 1
    result = n * fac(n - 1) % mod
    memo[n] = result
    return result

def solve(n, a, b):
    a = [[v - 1 for v in row] for row in a]
    b = [[v - 1 for v in row] for row in b]
    deltai = [-1] * n
    deltaj = [-1] * n

    def locs(g):
        l = [-1] * n ** 2
        for i in range(n):
            for j in range(n):
                assert l[g[i][j]] == -1
                l[g[i][j]] = (i, j)
        return l
    la = locs(a)
    lb = locs(b)
    found = [[False] * n for _ in range(n)]
    for i in range(n ** 2):
        (ai, aj) = la[i]
        (bi, bj) = lb[i]
        if ai != bi and aj != bj:
            di = (ai - bi) % n
            dj = (aj - bj) % n
            assert di and dj
            if found[di][dj]:
                return 0
            found[di][dj] = True
        if ai == bi:
            deltai[ai] = (bj - aj) % n
        if aj == bj:
            deltaj[aj] = (bi - ai) % n
    iseq = [[] for _ in range(n)]
    jseq = [[] for _ in range(n)]
    for i in range(n):
        for j in range(n):
            if deltai[i] > 0 and deltaj[j] > 0:
                if a[i][j - deltai[i]] != b[(i + deltaj[j]) % n][j]:
                    jseq[j].append(i)
                if a[i - deltaj[j]][j] != b[i][(j + deltai[i]) % n]:
                    iseq[i].append(j)
    jdeg = [0] * n
    for i in range(n):
        if deltai[i] > 0:
            for j in iseq[i]:
                if deltaj[j] > 0:
                    jdeg[j] += 1
    ideg = [0] * n
    for j in range(n):
        if deltaj[j] > 0:
            for i in jseq[j]:
                if deltai[i] > 0:
                    ideg[i] += 1
    ique = [i for i in range(n) if ideg[i] == 0 and deltai[i] > 0]
    jque = [j for j in range(n) if jdeg[j] == 0 and deltaj[j] > 0]
    ans = 1
    seq = []
    while ique or jque:
        ans = ans * fac(len(ique)) % mod
        for i in ique:
            seq.append(('i', i, deltai[i]))
            for j in iseq[i]:
                jdeg[j] -= 1
                if jdeg[j] == 0:
                    jque.append(j)
        ique = []
        ans = ans * fac(len(jque)) % mod
        for j in jque:
            seq.append(('j', j, deltaj[j]))
            for i in jseq[j]:
                ideg[i] -= 1
                if ideg[i] == 0:
                    ique.append(i)
        jque = []
    moved = [0] * n ** 2
    c = [[*row] for row in a]
    for (typ, x, d) in seq:
        assert 0 < d < n
        if typ == 'i':
            i = x
            vals = [c[i][j] for j in range(n)]
            vals = vals[-d:] + vals[:-d]
            for j in range(n):
                c[i][j] = vals[j]
        if typ == 'j':
            j = x
            vals = [c[i][j] for i in range(n)]
            vals = vals[-d:] + vals[:-d]
            for i in range(n):
                c[i][j] = vals[i]
        for v in vals:
            moved[v] += 1
        del typ, x, d, vals
    if any((v > 2 for v in moved)):
        return 0
    if c != b:
        return 0
    return ans

def main():

    def get_grid(n):
        return [[*map(int, input().split())] for _ in range(n)]
    for _ in range(int(input())):
        n = int(input())
        print(solve(n, get_grid(n), get_grid(n)) % mod)
if __name__ == '__main__':
    main()

Solution (Python):
import sys
mod = 998244353
(fact, ar0, xs, ys) = ([0] * 505, [0] * 505, [0] * 250005, [0] * 250005)
(ar, used) = ([], [])
t = int(input())
for xi in range(t):
    (ans, s, ar) = (1, 1, [])
    n = int(sys.stdin.readline())
    fact[0] = 1
    for i in range(1, n + 1):
        fact[i] = fact[i - 1] * i % mod
    for i in range(n):
        ar.append(sys.stdin.readline().split())
        for j in range(n):
            ar[i][j] = int(ar[i][j])
    for i in range(n):
        l = sys.stdin.readline().split()
        for j in range(n):
            xs[int(l[j])] = i
            ys[int(l[j])] = j
    used = [[0] * (n + 5) for i in range(n + 5)]
    for i in range(n):
        for j in range(n):
            (x0, y0) = ((xs[ar[i][j]] - i) % n, (ys[ar[i][j]] - j + n) % n)
            if x0 == 0 or y0 == 0:
                continue
            if used[x0][y0]:
                ans = 0
            used[x0][y0] = 1
    if ans == 0:
        print(0)
        continue
    (flag1, flag2) = ([0] * (n + 5), [0] * (n + 5))
    while s:
        s = 0
        for i in range(n):
            (tmp, fl) = (ys[ar[i][0]], 1)
            for j in range(1, n):
                fl &= tmp == (ys[ar[i][j]] - j) % n
            if fl == 0 or flag1[i] or tmp == 0:
                continue
            s += 1
            flag1[i] = 1
            for j in range(n):
                ar0[(j + tmp) % n] = ar[i][j]
            for j in range(n):
                ar[i][j] = ar0[j]
        ans = ans * fact[s] % mod
        s = 0
        for j in range(n):
            (tmp, fl) = (xs[ar[0][j]], 1)
            for i in range(1, n):
                fl &= tmp == (xs[ar[i][j]] - i) % n
            if fl == 0 or flag2[j] or tmp == 0:
                continue
            s += 1
            flag2[j] = 1
            for i in range(n):
                ar0[(i + tmp) % n] = ar[i][j]
            for i in range(n):
                ar[i][j] = ar0[i]
        ans = ans * fact[s] % mod
    for i in range(n):
        for j in range(n):
            if i == xs[ar[i][j]] and j == ys[ar[i][j]]:
                s += 1
    if s != n * n:
        ans = 0
    sys.stdout.write(str(ans) + '\n')

Solution (Python):
mod = 998244353
(fact, ar0, xs, ys) = ([0] * 505, [0] * 505, [0] * 250005, [0] * 250005)
(ar, used) = ([], [])
t = int(input())
for xi in range(t):
    (ans, s, ar) = (1, 1, [])
    n = int(input())
    fact[0] = 1
    for i in range(1, n + 1):
        fact[i] = fact[i - 1] * i % mod
    for i in range(n):
        ar.append(input().split())
        for j in range(n):
            ar[i][j] = int(ar[i][j])
    for i in range(n):
        l = input().split()
        for j in range(n):
            xs[int(l[j])] = i
            ys[int(l[j])] = j
    used = [[0] * (n + 5) for i in range(n + 5)]
    for i in range(n):
        for j in range(n):
            (x0, y0) = ((xs[ar[i][j]] - i) % n, (ys[ar[i][j]] - j + n) % n)
            if x0 == 0 or y0 == 0:
                continue
            if used[x0][y0]:
                ans = 0
            used[x0][y0] = 1
    if ans == 0:
        print(0)
        continue
    (flag1, flag2) = ([0] * (n + 5), [0] * (n + 5))
    while s:
        s = 0
        for i in range(n):
            (tmp, fl) = (ys[ar[i][0]], 1)
            for j in range(1, n):
                fl &= tmp == (ys[ar[i][j]] - j) % n
            if fl == 0 or flag1[i] or tmp == 0:
                continue
            s += 1
            flag1[i] = 1
            for j in range(n):
                ar0[(j + tmp) % n] = ar[i][j]
            for j in range(n):
                ar[i][j] = ar0[j]
        ans = ans * fact[s] % mod
        s = 0
        for j in range(n):
            (tmp, fl) = (xs[ar[0][j]], 1)
            for i in range(1, n):
                fl &= tmp == (xs[ar[i][j]] - i) % n
            if fl == 0 or flag2[j] or tmp == 0:
                continue
            s += 1
            flag2[j] = 1
            for i in range(n):
                ar0[(i + tmp) % n] = ar[i][j]
            for i in range(n):
                ar[i][j] = ar0[i]
        ans = ans * fact[s] % mod
    for i in range(n):
        for j in range(n):
            if i == xs[ar[i][j]] and j == ys[ar[i][j]]:
                s += 1
    if s != n * n:
        ans = 0
    print(ans)