yukicoder: No.1062 素敵なスコア

https://yukicoder.me/problems/no/1062

$A\lt B$ とします。
$K$ 以下の要素を黒丸、 $K$ よりも大きい要素を白丸で表すことにすると以下のような状況になります。

このうち操作を行った後も同じ要素になるindexは以下の黒枠で囲った部分になります。

このように $P,Q$ 両方でswapされずに同じ要素になるindexと、両方でswapされて同じ要素になるindexがあることがわかります。前者の総数は解説にあるように $(X^2+Y^2+Z^2) (N-1)!$ となります。
後者の総数を $S$ として、 $S$ を求めることを考えます。
各index $(1 \leq i \leq X)$ で $S_i$ =( $P_i,Q_i$ 両方で白丸で、操作で黒丸とswapされて $P_i=Q_i$ となる数列の総数)が求まれば、 $S=2\sum_{i=1}^X S_i$ となります。係数が2の理由ですが、○○がswapされる先は●●なので、 $P_i,Q_i$ 両方で黒丸のindexも同数だけあるからです。

$S_i$ の状況は以下のようになります。

これを、

こうじゃ。

なので、 $S_i$ =(後半に●●i個、前半に○●Y個と●●X-i個を好きなように並べた列に○○Z-1個を挿入した列の総数) $=X!Y!Z! \times {}_{Y+X-i}C_{Y}\times {}_{N-1}C_{Z-1}$ となります。よって、

$S=2X!Y!Z! \times {}_{N-1}C_{Z-1} \sum_{i=1}^X {}_{Y+X-i}C_{Y}=2X!Y!Z! \times {}_{N-1}C_{Z-1} \sum_{i=0}^{X-1} {}_{Y+i}C_{Y}$

ですが、経路数を少し考えると

$\sum_{i=0}^{X-1} {}_{Y+i}C_{Y}={}_{X+Y}C_{X-1}$

であることがわかるので、 $S=2X!Y!Z! \times {}_{N-1}C_{Z-1} \times {}_{X+Y}C_{X-1}=\frac{2XZ}{Y+1}(N-1)!$ となります。
前半パートと合わせれば答えは $(X^2+Y^2+Z^2+\frac{2XZ}{Y+1})(N-1)!$ です。

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const ll mod = 998244353;

int N, A, B;

int main() {
    cin >> N >> A >> B;
    if(A > B) swap(A, B);
    ll x = A, y = B - A, z = N - B;
    ll fac = 1, nfac = 1;
    for(int i = 1; i <= N - 1; i++) {
        fac = fac * i % mod;
        if(i != y + 1) nfac = nfac * i % mod;
    }
    cout << ((x * x % mod + y * y % mod + z * z % mod) * fac 
            + 2 * x * z % mod * nfac) % mod << endl;
}