\((A_1, \dots, A_N) < (A_{P_1}, \dots, A_{P_{N}})\) の条件を定義に従って場合分けします。

• \(A_1 < A_{P_1}\)
• \(A_1 = A_{P_1} \land A_2 < A_{P_2}\)
• \(A_1 = A_{P_1} \land A_2= A_{P_2} \land A_3 < A_{P_3}\)
• \(\vdots\)

それぞれの条件に付いて、条件を満たす \(A\) の個数を数えれば良いです。

例えば、\(A_1 < A_{P_1}\) を満たす \(A\) の個数は、\(P_1 = 1\) ならば \(0\) 、\(P_1 \neq 1\) ならば \(M^{N-2} \cdot \frac{M(M-1)}2\) です。

\(A_1 = A_{P_1} \land A_2 < A_{P_2}\) の場合は、\(A_1 = A_{P_1}\) なので、\(A_1\) と \(A_{P_1}\) を統合して同じ要素を指しているものと思えば (\(P_1 \neq 1\) ならば要素数が減る) 、\(A_2 < A_{P_2}\) のみが残り、上と同様に計算することができます。

したがって、Union-Find を利用して、\(i - P_i\) 間に辺を張っていくことで、これを高速に計算できます。

解答例 (C++)

#include <bits/stdc++.h>
#include <atcoder/modint>
#include <atcoder/dsu>
using namespace std;
using mint = atcoder::modint998244353;

int main(){
    int N, M;
    cin >> N >> M;
    vector<int> P(N);
    for(int& p : P){
        cin >> p;
        p--;
    }
    int n = N;
    mint ans = 0;
    atcoder::dsu uf(N);
    for(int i = 0; i < N; i++){
        const int p = P[i];
        if(uf.same(i, p)) continue;
        ans += mint{M}.pow(n - 1) * (M - 1) / 2;
        n--;
        uf.merge(i, p);
    }
    cout << ans.val() << endl;
}