つちのこの日記

競プロと音ゲーがメインです。

SolveMe (AOJ 2378)

この記事は帰ってきた AOJ-ICPC Advent Calendar 2022 25 日目の記事です。

問題リンク

https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2378 (1200+ 点)

問題概要

集合 $S={0,1,\ldots,N-1}$ にたいして、 $S\to S$ の写像 $f,g$ を考える。 $(f,g)$ の組であって、$f ^ X g f ^Y g f ^Z$ が恒等写像となるものは何通りあるか、$1000000007$ で割ったあまりを求めよ。

制約

  • $1 \leq N_A,N_B \leq 1000$
  • $ |X|,|Y|,|Z|\leq 10 ^ {18}$
  • 入力は全て整数

解法

$X=Y=Z=0$ のときはコーナーケースなので別途処理します。そうでない場合、$f,g$ は全単射であることが要請されます。式変形により、

$$(f ^ Y g) ^ 2 = f ^{(Y-X-Z)} $$

を満たす $(f,g)$ を数え上げればよいです。適切に置きなおし、$Y-X-Z<0$ のときは逆写像を考えることで、条件は

$$f ^ 2 = g ^{K}\ (K = |Y-X-Z|) $$

と表せます。

直接 $(f,g)$ を数え上げるのは難しいので、一工夫します。$h=f ^2 (=g ^ K)$ と置き、$h$ すべてに対して、($f ^ 2=h$ なる $f$ の個数) $\times$ ($g ^K=h$ なる $f$ の個数) を求め、足し合わせたものが答えと一致するので、そちらを数え上げます。

まずは、$h$ についてです。写像 $f$ はいくつかのサイクルに分解されますが、これを $k$ 乗すると、長さ $N$ のサイクル $1$ 個は 長さ $N/\mathrm{gcd}(k,N)$ のサイクル $\mathrm{gcd}(k,N)$ 個になることがわかります。

なので、まず $h$ に含まれる長さ $i$ のサイクルの個数を決め打ち、次にそれを生成する $f,g$ の個数を掛け合わせていけばよいです。前者は単純な動的計画法で出来ます。

後者についてです。長さ $i$ のサイクルが $j$ 個あるとき、それを生成する方法を求めればよいですが、これも同様の動的計画法で出来ます。

ただし、複数のサイクルを組み合わせて一つのサイクルを作るとき、サイクルを組み合わせる順番、組み合わせる際のオフセット、などやや丁寧に考えるべき点が多いです。自分はオフセットの考慮を忘れていて無限に困っていました。

サンプルが弱いこともあり、デバッグがそのままでは大変しづらいので、愚直テストを書くのが有効な問題だと思いました。

実装例 (c++)

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using pii = pair<int, int>;
using vi = vector<int>;
#define rep(i, x)    for(int i = 0; i < (x); i++)
#define rng(i, l, r) for(int i = (l); i < (r); i++)
#define sz(v)        (int)(v).size()
#define all(v)       (v).begin(), (v).end()

constexpr int mod = 1000000007;
struct mint {
    ll v;
    mint(ll x = 0) {
        v = x;
        v %= mod;
        if(v < 0) v += mod;
    }
    mint operator+(mint o) { return v + o.v; }
    mint operator-(mint o) { return v - o.v; }
    mint operator*(mint o) { return v * o.v; }
    mint pow(ll k) {
        assert(k >= 0);
        mint ret = 1, a = v;
        while(k) {
            if(k & 1) ret = ret * a;
            a = a * a;
            k >>= 1;
        }
        return ret;
    }
    mint inv() { return pow(mod - 2); }
    mint operator/(mint o) { return v * o.inv().v; }
};

vector<mint> fac, finv;

void init(int n) {
    fac.resize(n + 1);
    finv.resize(n + 1);
    fac[0] = finv[0] = 1;
    rep(i, n) fac[i + 1] = fac[i] * (i + 1);
    finv[n] = fac[n].inv();
    for(int i = n - 1; i >= 1; i--) finv[i] = finv[i + 1] * (i + 1);
}
mint binom(int n, int k) {
    if(min(n, k) < 0 or n < k) return 0;
    return fac[n] * finv[k] * finv[n - k];
}

int main() {
    int n;
    ll X, Y, Z;
    cin >> n >> X >> Y >> Z;

    if(X == 0 and Y == 0 and Z == 0) {
        vector<mint> a{1, 1};
        rng(i, 2, n + 1) a.push_back(a[i - 1] + a[i - 2] * (i - 1));
        cout << (a[n] * mint(n).pow(n)).v << endl;
        return 0;
    }

    init(1000);
    ll k = abs(Y - X - Z);

    //  (i = j / gcd(j, k)) * gcd(j, k)
    auto make_table = [&](ll x) {
        vector table(n + 1, vector(n + 1, mint(0)));
        rng(i, 1, n + 1) {
            vi tmp;
            rng(j, 1, n + 1) {
                if(j / gcd(j, x) != i) continue;
                tmp.push_back(gcd(j, x));
            }
            if(tmp.empty()) continue;
            int m = sz(tmp);
            // j個のグループを、tmp[k]以下を使って分配する個数
            vector dp(m + 1, vector(n / i + 1, mint(0)));
            dp[0][0] = 1;
            rep(id, m) {
                mint pw = mint(i).pow(tmp[id] - 1);
                rep(j, n / i + 1) {
                    mint tar = 1;
                    for(int u = 0; u * tmp[id] <= j; u++) {
                        dp[id + 1][j] = dp[id + 1][j] + dp[id][j - u * tmp[id]] * tar * finv[u];
                        tar = tar * binom(j - u * tmp[id], tmp[id]) * fac[tmp[id] - 1] * pw;
                    }
                }
            }
            rep(j, n / i + 1) table[i][j] = dp.back()[j];
        }
        return table;
    };

    auto lt = make_table(2);
    auto rt = make_table(k);


    vector dp(n + 1, vector(n + 1, mint(0)));
    dp[0][n] = 1;
    rng(i, 1, n + 1) {
        rep(j, n + 1) {
            dp[i][j] = dp[i][j] + dp[i - 1][j];
            //完成形のグループを決め打つ
            //それを適切に分配する
            mint cur = 1;
            for(int u = 1; u * i <= j; u++) {
                cur = cur * binom(j - i * (u - 1), i) * fac[i - 1];
                mint tar = cur / fac[u];
                mint lef = lt[i][u], rig = rt[i][u];
                dp[i][j - u * i] = dp[i][j - u * i] + dp[i - 1][j] * tar * lef * rig;
            }
        }
    }
    cout << dp[n][0].v << endl;
}

感想

近年の数え上げのインフレ?もあり、1200+ にしてはかなり解きやすいように感じました。ただ、正確な数え上げを要求されることもあり、中々芯のある問題で面白かったです。