B - Counting Grids AtCoder Regular Contest 143
N×N のマス目の各マスに 1 から N2 までの整数を 1 つずつ書き込む方法であって， どのマスも以下の条件のうち少なくとも一方を満たすようなものの個数を998244353 で割ったあまりを求めてください．条件A : そのマスに書かれている数より大きい数が同じ列に存在する条件B : そのマスに書かれている数より小さい数が同じ行に存在する
N×Nの各マスに、1~Nの数字１つずつを、AまたはBの条件を満たしながら、どんどん書き込みます。それで、その書き込み方の個数を求めます。
考え方の方針
直接「${A\cup B}$を満たすマス」を考えるのはしんどいので、その逆、「${\bar{A} \cap \bar{B}}$を満たすマス」を考えれば、問題を「すべてのマスが${\bar{A}\cap \bar{B}}$を満たす書き込み方の個数を求める」ことに置き換えられそうです。
${\bar{A}\cap \bar{B}}$
条件${\bar{A}}$ : そのマスに書かれている数より小さい数が同じ列に存在する、かつ
条件${\bar{B}}$: そのマスに書かれている数より大きい数が同じ行に存在する
数式で書くと
\begin{array}{l}n(A\cup B) &=& n(U) - n(\overline{A\cup B}) \\ &=& n(U) - n(\bar{A}\cap \bar{B}) \end{array}
こんな感じです。
それでは実際に、${n(\bar{A}\cap \bar{B})}$がいくつになるかを計算してみましょう。
計算の手順
（１） １つのマスに注目する
１つのマスに注目する方法は全部で
${N^2}$通り${\cdots ①}$
あります。
（２） 行・列の数の組み合わせ
先ほど注目したマスを含む、行と列を考えます。
ここで、水色のマスは$2N - 1$あるので、そこに数を書き込む組み合わせは全部で、
${}_{N^2} \mathrm{C} _{2N-1}$通り${\cdots②}$
あります。
（３） 条件を満たす数を書き込む
さらに、この水色のマスに書き込む数は、${\bar{A}\cap \bar{B}}$を満たす必要があるわけですが、その書き込み方を考えてみましょう。まずは、注目したマス（黄緑のマス）には何が入るでしょうか。実は１通りしかありません。
例えば、N=3で、（２）で選ばれた数の組み合わせが{2, 3, 6, 7, 9}だったとしましょう。条件を満たすには、中央値が黄緑色のマスに入らなければならないとわかります。このケースの場合は「６」が中央値です。
残りの水色のマスには、{2, 3}と{7, 9}の並べ方で決まるので、全部で２!×２! = ４通りあります。
最後に、空白の部分には残りの数を自由に書き込みます。全部で４！= ２４通りです。
よって、一般にNを使って書き込み方の個数を表すと、
${(N-1)!(N-1)!{(N-1)^2}!}$ 通り${\cdots③}$
になります。N=3のときは、
${2!\times2!\times4!=97}$通りです。
（４） ①②③を使って個数を計算
①, ②, ③より
$${n(\bar{A} \cap \bar{B}) = n^2 \times _{N^2} \mathrm{C} _{2N-1}\times (N-1)!(N-1)!(N-1)^2! }$$
となり、全体集合をUとすると${n(U)=N^2!}$なので
{\begin{array}{l} n(A\cup B) &=& n(U) - n(\bar{A}\cap \bar{B}) \\ &=& N^2! - N^2 \times _{N^2} \mathrm{ C } _{2N-1}\times (N-1)!(N-1)!(N-1)^2! \end{array}}
ですべての書き込み方の個数が求まりました。
Pythonで実装
階乗の計算は時間がかかるので、メモ化しておきましょう。
import math, functools
@functools.lru_cache(maxsize=None)
def fac(n):
return math.factorial(n)
def nCr(n, r):
return fac(n) // (fac(r) * fac(n-r))
N = int(input())
M = 998244353
U = fac(N*N)
a = N**2
b = nCr(N**2, 2*N - 1)
c = (fac(N-1)**2) * (fac((N-1)**2))
print((U - a*b*c)%M)
以上です。何か間違ってたらコメントくれると嬉しいです。