Official

Ex - Dice Sum Infinity Editorial by MMNMM


これまでに出た目の合計 \(X\) が \(R-X\equiv i\pmod{10^9}\) を満たす状態から、操作を終了するまでにサイコロを振る回数の期待値 \(E _ i\) を求めることを考えます。

\(X-R\geq k10^9\) となるまでに、必ず \(X-R=k10^9-1,k10^9-2,k10^9-3,k10^9-4,k10^9-5,k10^9-6\) のいずれかの状態になることから、\(E _ 1,E _ 2,E _ 3,E _ 4,E _ 5,E _ 6\) を \(E _ 1,E _ 2,E _ 3,E _ 4,E _ 5,E _ 6\) のみの式で書くことができます。

次のような問題を考えます。

問題

整数 \(r\) が与えられる。 サイコロを振って出た目の数を \(r\) から引くことを繰り返す。 \(r\leq 0\) となったら操作を終える。 操作を終えるまでにサイコロを振る回数の期待値 \(e( r )\) と、操作を終えた時点での \(r\) の値が \(i\) と等しい確率 \(p _ i(r )\ (i=0,-1,-2,-3,-4,-5)\) を求めよ。

これらは線形漸化式を持つので、\(r\) が \(1\) つ与えられるごとに \(O(\log r)\) 時間で値を求めることができます(サイコロの面が \(d\) 個あるとしたとき、\(O(d^2\log d\log r)\) 時間などで解くことができます)。

よって、\(E _ i\ (i=1,2,3,4,5,6)\) の満たす等式が \(6\) 本得られ、これを解くと \(E _ i\) がわかります。 時間計算量は \(O(d^3\log d\log r)\) などになります。 入力によらないため、事前に計算しておくこともできます。

この問題の答えは \(\displaystyle e(R-6)+\sum _ {i=1} ^ 6E _ ip _ {i-6}(R-6)\) です。

よって、この問題を \(O(\log r)\) 時間で解くことができました。

また、この答えは線形漸化式を持つ数列の定数倍の和なので、線形漸化式を持ちます。 線形漸化式を事前に求めておくことで、計算量を減らすことができます。

実装例は以下のようになります。

式変形を執拗にした版

#include <iostream>
#include <atcoder/modint>

int main() {
    using namespace std;
    using modint = atcoder::static_modint<998244353>;

    unsigned r;
    cin >> r;
    modint ans{modint::raw(335634206).pow(r + 871641306) +
               modint::raw(920597818).pow(r + 227058393) +
               modint::raw(674695235).pow(r + 903556278) +
               modint::raw(713031681) * r +
               modint::raw(488866572)};
    pair<modint, modint> num{244113007, 739484271}, den{766308848, 522180481};
    auto&&[a, b]{num};
    auto&&[c, d]{den};
    while (r) {
        if (r & 1) num = {b - a * c, b * d};
        else num = {a, a * d - b * c};
        den = {2 * d - c * c, d * d};
        r /= 2;
    }
    cout << (ans + num.first).val() << endl;
    return 0;
}

式変形をあまりしない版

#include <bits/extc++.h>
#include <atcoder/modint>

template<int m>
std::ostream &operator<<(std::ostream &os, const atcoder::static_modint<m> &x) {
    os << x.val();
    return os;
}

int main() {
    using namespace std;
    using modint = atcoder::static_modint<998244353>;

    // [x^r]f(x)/g(x) を求める
    auto calc_linear{[poly_mul{[](const auto &lhs, const auto &rhs, unsigned long parity = 0) {
        const auto L{lhs.size()};
        const auto R{rhs.size()};
        vector<modint> ret((L + R - parity + 1) / 2);
        for (unsigned long i{}; i < L; ++i)
            for (unsigned long j{(parity ^ i) & 1}; j < R; j += 2)
                ret[(i + j) / 2] += lhs[i] * rhs[j];
        while (!ret.empty() && ret.back() == 0)ret.pop_back();
        return ret;
    }}](auto num, auto den, auto r) mutable -> modint {
        vector <modint> den_minus{den};
        while (r) {
            copy(begin(den), end(den), begin(den_minus));
            for (unsigned long i{1}; i < den_minus.size(); i += 2)den_minus[i] *= -1;
            num = poly_mul(num, den_minus, r & 1);
            den = poly_mul(den, den_minus);
            r /= 2;
        }
        return num[0] / den[0];
    }};

    // 部分問題の r に対する答えを求める
    const auto far_to_near_coeffs{[&calc_linear](unsigned long r) {
        vector<modint> c(6);
        modint x;
        const modint d{modint::raw(6).inv()};
        vector<modint> den{1, -d * 7, 0, 0, 0, 0, 0, d};
        c[0] = calc_linear(vector<modint>{1, -1}, den, r);
        c[1] = calc_linear(vector<modint>{0, d, 0, 0, 0, 0, -d}, den, r);
        c[2] = calc_linear(vector<modint>{0, d, 0, 0, 0, -d}, den, r);
        c[3] = calc_linear(vector<modint>{0, d, 0, 0, -d}, den, r);
        c[4] = calc_linear(vector<modint>{0, d, 0, -d}, den, r);
        c[5] = calc_linear(vector<modint>{0, d, -d}, den, r);
        x = calc_linear(vector<modint>{0, 1}, den, r);
        return make_pair(c, x);
    }};

    // E_i を求める
    vector<modint> expected_value(6);
    {
        vector<pair<vector<modint>, modint>> coefs;
        for (unsigned long i{7}; i < 12; ++i) {
            const auto&[a, b]{far_to_near_coeffs(1000000000 - i)};
            coefs.emplace_back(a, b);
        }
        // E_i が満たす式 6 本を計算する
        vector matrix(6, vector<modint>(7));
        for (unsigned long i{}; i < 6; ++i) {
            matrix[i][i] += 1;
            matrix[i][6] += 1;
            for (unsigned long j{1}; j <= 6; ++j)
                if (i + j < 6)
                    matrix[i][i + j] -= modint{1} / modint{6};
                else if (i + j > 6) {
                    for (unsigned long k{}; k < 6; ++k)
                        matrix[i][k] -= coefs[i + j - 7].first[k] / 6;
                    matrix[i][6] += coefs[i + j - 7].second / 6;
                }
        }

        // E_i を求める
        for (unsigned long i{}; i < 6; ++i) {
            {
                const modint g{modint{1} / matrix[i][i]};
                for (unsigned long j{}; j <= 6; ++j)matrix[i][j] *= g;
            }
            for (unsigned long j{}; j < 6; ++j)
                if (j != i) {
                    const modint c{matrix[j][i]};
                    for (unsigned long k{}; k <= 6; ++k)
                        matrix[j][k] -= matrix[i][k] * c;
                }
        }
        
        for (unsigned long i{}; i < 6; ++i)expected_value[i] = matrix[i].back();
    }


    unsigned long r;
    cin >> r;
    if (r <= 6) {
        cout << expected_value[6 - r] << endl;
    } else {
        const auto&[probability, expected]{far_to_near_coeffs(r - 6)};
        cout << inner_product(begin(probability), end(probability), begin(expected_value), expected) << endl;
    }
    return 0;
}

posted:
last update: