G - Sum of Pow of Mod of Linear 解説 by MMNMM


Floor Sum のモノイド積版を用いることでこの問題を解くことができます。

Floor Sum のモノイド積版については、丸紅プログラミングコンテスト2024 F の maspy さんによるユーザ解説や、MMNMM のブログ記事を参照してください。


Floor Sum のモノイド積版によって、非負整数 \(a,b\) と正整数 \(n,m\) および可換環の元 \(A,B\) に対して、\(\displaystyle\sum _ {i=0} ^ {n-1}A ^ iB ^ {\lfloor(ai+b)/m\rfloor}\) の値を \(O(\log n+\log m+\log a+\log b)\) 時間で求めることができます(可換でなくてもできますが、可換性を利用すると少しだけ簡単になります)。

\(X\) が\({}\bmod R\) で逆元をもつなら、求めるものは \(\displaystyle X ^ B\sum _ {i=0} ^ {N-1}(X ^ A) ^ i(X ^ {-M}) ^ {\lfloor(Ai+B)/M\rfloor}\) となり、Floor Sum のモノイド積版を用いて値が得られます。

\(R\) を、\(\gcd(R _ 0,X)=1\) かつ十分大きな \(n\) に対して \(R _ 1| X ^ n\) が成り立つような \(R=R _ 0R _ 1\) に分解することを考えます(\(R\) の素因数のべきそれぞれに対して、\(X\) もその素因数をもてば \(R _ 1\) に、そうでなければ \(R _ 0\) にその素べきをかけることで構成的に得られます)。 中国剰余定理を使うことで、\({}\bmod R _ 0\) での答えと\({}\bmod R _ 1\) での答えを得られれば\({}\bmod R\) での答えが得られます。

\(\gcd(R _ 0,X)=1\) なので \(X\) は\({}\bmod R _ 0\) で逆元を持ちます。このことから\({}\bmod R _ 0\) での答えは上述のアルゴリズムで得られます。

\(n _ 0\le n\) なる任意の \(n\) に対して \(R _ 1|X ^ n\) が成り立つような \(n _ 0\) に対して、\(\!{}\bmod R _ 1\) での答えは

\[\sum _ {0\le i\lt N,0\le Ai+B\bmod M\lt n _ 0}X ^ {Ai+B\bmod M}\pmod{R _ 1}\]

となります。 適切な \(n _ 0\) をとり、\(Ai+B\bmod M\) の値で分類して主客転倒をすることで、\(O(n _ 0)=O(\log R)\) 時間でこの値を求めることができます。

よって、全体で \(1\) ケースあたり \(O(\log R+\log A+\log B+\log M+\log N)\) 時間となり、十分高速です。

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

// モノイドのべき乗
template <typename Monoid, typename I>
auto monoid_pow(Monoid x, I exp) {
    Monoid res{Monoid::unit()};

    while (exp) {
        if (exp & 1)
            res *= x;

        x *= x;
        exp >>= 1;
    }

    return res;
}

// floor sum のモノイド積版
template <typename I, class Monoid>
Monoid floor_prod(I n, I m, I a, I b, Monoid x, Monoid y) {
    x *= monoid_pow(y, a / m);
    a %= m;

    const auto prefix_prod{monoid_pow(y, b / m)};
    b %= m;

    const auto y_max{a * n + b};

    if (y_max < m)
        return prefix_prod * monoid_pow(x, n);

    return prefix_prod * floor_prod(y_max / m - 1, a, m, m + a - b - 1, y, x) * y * monoid_pow(x, y_max % m / a);
}

// 可換環の元 A, B に対して x := (A, 1), y := (B, 0) として floor sum のモノイド積版をすると ∑ A^i B^floor(Ai+B/M) が得られる
template <class T>
class commutative_ring_power_sum {
public:
    using value_type = T;
    T a, s;
    commutative_ring_power_sum(T _a, T _s) : a(_a), s(_s) {}

    static commutative_ring_power_sum X(T _a) {
        return commutative_ring_power_sum{_a, T(1)};
    }

    static commutative_ring_power_sum Y(T _b) {
        return commutative_ring_power_sum{_b, T()};
    }

    static commutative_ring_power_sum unit() {
        return commutative_ring_power_sum{T(1), T()};
    }

    commutative_ring_power_sum& operator*=(commutative_ring_power_sum const& rhs) {
        s += a * rhs.s;
        a *= rhs.a;
        return *this;
    }

    commutative_ring_power_sum operator*(commutative_ring_power_sum const& other) const {
        return commutative_ring_power_sum(*this) *= other;
    }

    const T& ans() const {
        return s;
    };
};

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

int main() {
    using namespace std;
    using modint = atcoder::modint;
    unsigned T;
    cin >> T;
    for (unsigned t{}; t < T; ++t) [] {
        unsigned N, M, A, B, X, R;
        cin >> N >> M >> A >> B >> X >> R;
        unsigned mod{R};
        for (unsigned i{}; i < 30; ++i)
            mod /= gcd(mod, X);
        modint::set_mod(mod);
        // 逆元があるほう
        const auto v1{(floor_prod<unsigned long>(N, M, A, B, commutative_ring_power_sum<modint>::X(modint{X}.pow(A)), commutative_ring_power_sum<modint>::Y(modint{X}.pow(M).inv())).ans() * modint{X}.pow(B)).val()};
        // 冪零のほう
        const auto v2{[rem_mod{R / mod}, N, M, A, B, X] {
            unsigned ans{};
            const unsigned long d{gcd(M, A)}, a{A / d}, m{M / d}, coef(atcoder::inv_mod(a, m));
            unsigned long tmp{1};
            for (unsigned i{}; i < M; ++i) {
                if (const auto first{(i + M - B) / d * coef % m}; i % d == B % d && first < N)
                    ans += (N + m - first - 1) / m * tmp % rem_mod;
                (tmp *= X) %= rem_mod;
                if (tmp == 0)
                    break;
            }
            return ans % rem_mod;
        }()};
        cout << atcoder::crt({v1, v2}, {mod, R / mod}).first << endl;
    }();
    return 0;
}

投稿日時:
最終更新: