確かに、えでゅふぉのラス問にありそう！

問題概要

$N$ 人の選手が距離 $L$ のレースを走る。 $i (= 1, 2, \dots, N)$ 人目はデフォルトでは距離 1 走るのに $t_{i}$ 秒かかる。

スタートから距離が $0 \le X_{1} \lt X_{2} \lt \dots \lt X_{M} \lt L$ のところにアイテムがある。各アイテムは各選手に対して独立に、確率 $\frac{p_{i}}{100}$ で減速 (距離 1 走るのにかかる時間が $t_{i} + D$ 秒になる) し、確率 $\frac{100 - p_{i}}{100}$ で加速 (距離 1 走るのにかかる時間が $t_{i} - D$ 秒になる) する。

1 番目の選手の順位の期待値を求めよ。ただし、選手 $i, j$ が同着のとき、 $i \lt j$ ならば選手 $i$ の方が順位が上とする。

制約

$2 \le N, L, t_{i} \le 10^{5}$
$1 \le M \le L - 1$

考えたこと

すべて 0-indexed で考えることにする。選手 $0$ の順位が $a$ であるとは、選手 $1, 2, \dots, N-1$ のうち選手 $0$ に勝った人が $a - 1$ 人であることを意味する。よって求める期待値は、期待値の線形性から、

$1 + \sum_{i = 2}^{N} (選手 i が選手 1 に勝つ確率)$

ということになる。選手 $i$ が選手 $0$ に勝つ確率を求めるためには、次のものを求める必要がある。

選手 $0$ がゴールするまでの時間の確率分布を求める
選手 $i$ がゴールするまでの時間の確率分布を求める

これらは自体は独立に計算できる。

距離 1 あたり t 秒で走る選手のゴール時間の確率分布

選手の速度を $t$ [ 秒/m] とする。選手が何秒でゴールするかについては $2^{M}$ 通りのパターンがある。しかしながら、減速する分の距離 $l$ (このとき、加速する距離は $L - X_{0} - l$ になる) としてありうるパターンは $O(L)$ 通りになるので、それほど多くはない。加速する分の距離が $l$ であるとき、選手のゴール時間は

$Lt + lD - (L - X_{0} - l)D$

となる。というわけで、選手が減速分距離が $l$ となる確率を計算しよう。普通に考えれば

${\rm dp} \lbrack i \rbrack \lbrack l \rbrack$ := 最初の $i$ 個のチェックポイントを使った時点で、減速距離が $l$ になる確率

として、部分和 DP をすればよい。でもこれでは $O(ML)$ の計算量を要するので間に合わない。こういうのを FFT で高速化するのは典型ではある。

部分和 DP を求める問題を FFT へ

結局、各チェックポイントごとに加速または減速する距離を $l_{0}, l_{1}, \dots, l_{M}$ として、次の多項式を考えたとき、 $x^{l}$ 次の項が求める確率となる！！！

$(p_{0}x^{l_{0}} + (1 - p_{0}))(p_{1}x^{l_{1}} + (1 - p_{1})) \dots (p_{M-1}x^{l_{M-1}} + (1 - p_{M-1}))$

これも単純に FFT すると、 $O(L^{2} \log L)$ となってしまう。しかし順番を工夫すると $O(L (\log L)^{2})$ になるのだ。具体的には次のようにすればよい。

$M$ 個ある多項式のうち次数の小さい 2 個を選んで計算し、その積で置き換えることを繰り返す

このように順序を工夫する発想で計算量を下げる考え方は、次の問題でも出ていた模様。

codeforces.com

確率分布から答えへ

各選手の減速距離に関する確率分布が求められたら、選手 $i$ が選手 $0$ に勝つ確率は、次のように計算できる。

選手 $0$ の減速距離が $l_{0}$ であるとき
選手 $i$ の減速距離を $l_{i}$ として

$Lt_{0} + l_{0}D - (L - X_{0} - l_{0})D \gt Lt_{i} + l_{i}D - (L - X_{0} - l_{i})D$
⇔ $l_{i} - l_{0} \lt \frac{L(t_{0} - t_{i})}{2D}$

となる。つまり、 $l_{i} - l_{0}$ の確率分布がわかれば、次のようにして答えが求められる。

$l_{i} - l_{0}$ の確率分布の累積和をとる
それを用いて $\frac{L(t_{0} - t_{i})}{2D}$ 未満となる確率を求める

$l_{i} - l_{0}$ の確率分布は、ふたたび FFT を用いて計算できる。計算量は $O(L (\log L)^{2} + N)$ 。

コード

#include <bits/stdc++.h>
using namespace std;
template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; }
template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return 1; } return 0; }

// modint
template<int MOD> struct Fp {
    long long val;
    constexpr Fp(long long v = 0) noexcept : val(v % MOD) {
        if (val < 0) val += MOD;
    }
    constexpr int getmod() const { return MOD; }
    constexpr Fp operator - () const noexcept {
        return val ? MOD - val : 0;
    }
    constexpr Fp operator + (const Fp& r) const noexcept { return Fp(*this) += r; }
    constexpr Fp operator - (const Fp& r) const noexcept { return Fp(*this) -= r; }
    constexpr Fp operator * (const Fp& r) const noexcept { return Fp(*this) *= r; }
    constexpr Fp operator / (const Fp& r) const noexcept { return Fp(*this) /= r; }
    constexpr Fp& operator += (const Fp& r) noexcept {
        val += r.val;
        if (val >= MOD) val -= MOD;
        return *this;
    }
    constexpr Fp& operator -= (const Fp& r) noexcept {
        val -= r.val;
        if (val < 0) val += MOD;
        return *this;
    }
    constexpr Fp& operator *= (const Fp& r) noexcept {
        val = val * r.val % MOD;
        return *this;
    }
    constexpr Fp& operator /= (const Fp& r) noexcept {
        long long a = r.val, b = MOD, u = 1, v = 0;
        while (b) {
            long long t = a / b;
            a -= t * b, swap(a, b);
            u -= t * v, swap(u, v);
        }
        val = val * u % MOD;
        if (val < 0) val += MOD;
        return *this;
    }
    constexpr bool operator == (const Fp& r) const noexcept {
        return this->val == r.val;
    }
    constexpr bool operator != (const Fp& r) const noexcept {
        return this->val != r.val;
    }
    constexpr bool operator < (const Fp& r) const noexcept {
        return this->val < r.val;
    }
    friend constexpr istream& operator >> (istream &is, Fp<MOD>& x) noexcept {
        is >> x.val;
        x.val %= MOD;
        if (x.val < 0) x.val += MOD;
        return is;
    }
    friend constexpr ostream& operator << (ostream &os, const Fp<MOD>& x) noexcept {
        return os << x.val;
    }
    friend constexpr Fp<MOD> modpow(const Fp<MOD> &a, long long n) noexcept {
        if (n == 0) return 1;
        auto t = modpow(a, n / 2);
        t = t * t;
        if (n & 1) t = t * a;
        return t;
    }
};

namespace NTT {
    long long modpow(long long a, long long n, int mod) {
        long long res = 1;
        while (n > 0) {
            if (n & 1) res = res * a % mod;
            a = a * a % mod;
            n >>= 1;
        }
        return res;
    }

    long long modinv(long long a, int mod) {
        long long b = mod, u = 1, v = 0;
        while (b) {
            long long t = a / b;
            a -= t * b, swap(a, b);
            u -= t * v, swap(u, v);
        }
        u %= mod;
        if (u < 0) u += mod;
        return u;
    }

    int calc_primitive_root(int mod) {
        if (mod == 2) return 1;
        if (mod == 167772161) return 3;
        if (mod == 469762049) return 3;
        if (mod == 754974721) return 11;
        if (mod == 998244353) return 3;
        int divs[20] = {};
        divs[0] = 2;
        int cnt = 1;
        long long x = (mod - 1) / 2;
        while (x % 2 == 0) x /= 2;
        for (long long i = 3; i * i <= x; i += 2) {
            if (x % i == 0) {
                divs[cnt++] = i;
                while (x % i == 0) x /= i;
            }
        }
        if (x > 1) divs[cnt++] = x;
        for (int g = 2;; g++) {
            bool ok = true;
            for (int i = 0; i < cnt; i++) {
                if (modpow(g, (mod - 1) / divs[i], mod) == 1) {
                    ok = false;
                    break;
                }
            }
            if (ok) return g;
        }
    }

    int get_fft_size(int N, int M) {
        int size_a = 1, size_b = 1;
        while (size_a < N) size_a <<= 1;
        while (size_b < M) size_b <<= 1;
        return max(size_a, size_b) << 1;
    }

    // number-theoretic transform
    template<class mint> void trans(vector<mint> &v, bool inv = false) {
        if (v.empty()) return;
        int N = (int)v.size();
        int MOD = v[0].getmod();
        int PR = calc_primitive_root(MOD);
        static bool first = true;
        static vector<long long> vbw(30), vibw(30);
        if (first) {
            first = false;
            for (int k = 0; k < 30; ++k) {
                vbw[k] = modpow(PR, (MOD - 1) >> (k + 1), MOD);
                vibw[k] = modinv(vbw[k], MOD);
            }
        }
        for (int i = 0, j = 1; j < N - 1; j++) {
            for (int k = N >> 1; k > (i ^= k); k >>= 1);
            if (i > j) swap(v[i], v[j]);
        }
        for (int k = 0, t = 2; t <= N; ++k, t <<= 1) {
            long long bw = vbw[k];
            if (inv) bw = vibw[k];
            for (int i = 0; i < N; i += t) {
                mint w = 1;
                for (int j = 0; j < t/2; ++j) {
                    int j1 = i + j, j2 = i + j + t/2;
                    mint c1 = v[j1], c2 = v[j2] * w;
                    v[j1] = c1 + c2;
                    v[j2] = c1 - c2;
                    w *= bw;
                }
            }
        }
        if (inv) {
            long long invN = modinv(N, MOD);
            for (int i = 0; i < N; ++i) v[i] = v[i] * invN;
        }
    }

    // small case (T = mint, long long)
    template<class T> vector<T> naive_mul 
    (const vector<T> &A, const vector<T> &B) {
        if (A.empty() || B.empty()) return {};
        int N = (int)A.size(), M = (int)B.size();
        vector<T> res(N + M - 1);
        for (int i = 0; i < N; ++i)
            for (int j = 0; j < M; ++j)
                res[i + j] += A[i] * B[j];
        return res;
    }

    // mint
    template<class mint> vector<mint> operator * 
    (const vector<mint> &A, const vector<mint> &B) {
        if (A.empty() || B.empty()) return {};
        int N = (int)A.size(), M = (int)B.size();
        if (min(N, M) < 30) return naive_mul(A, B);
        int MOD = A[0].getmod();
        int size_fft = get_fft_size(N, M);
        vector<mint> a(size_fft), b(size_fft), c(size_fft);
        for (int i = 0; i < N; ++i) a[i] = A[i];
        for (int i = 0; i < M; ++i) b[i] = B[i];
        trans(a), trans(b);
        vector<mint> res(size_fft);
        for (int i = 0; i < size_fft; ++i) res[i] = a[i] * b[i];
        trans(res, true);
        res.resize(N + M - 1);
        return res;
    }
};

const int MOD = 998244353;
using mint = Fp<MOD>;
using namespace NTT;

int main() {
    long long N, L, M, D;
    cin >> N >> L;
    vector<long long> t(N);
    for (int i = 0; i < N; ++i) cin >> t[i];
    cin >> M >> D;
    vector<long long> x(M + 1, L);
    vector<mint> p(M);
    for (int i = 0; i < M; ++i) cin >> x[i] >> p[i], p[i] /= 100;

    priority_queue<pair<int,vector<mint>>, vector<pair<int,vector<mint>>>, greater<pair<int,vector<mint>>>> que;
    for (int i = 0; i < M; ++i) {
        long long l = x[i+1] - x[i];
        vector<mint> v(l+1, 0);
        v[0] = mint(1) - p[i];
        v[l] = p[i];
        que.push({l+1, v});
    }
    while (que.size() >= 2) {
        auto left = que.top(); que.pop();
        auto right = que.top(); que.pop();
        auto mul = left.second * right.second;
        que.push({mul.size(), mul});
    }
    auto f = que.top().second;
    long long deg = f.size() - 1;
    auto g = f;
    reverse(g.begin(), g.end());
    auto fg = f * g;
    vector<mint> sum(fg.size()+1, 0);
    for (int i = 0; i < fg.size(); ++i) sum[i+1] = sum[i] + fg[i];

    mint res = 1;
    for (int i = 1; i < N; ++i) {
        long long lim = deg + (L * (t[0] - t[i]) + D * 20000000 + D * 2 - 1) / (D * 2) - 10000000;
        chmax(lim, 0LL);
        chmin(lim, (long long)sum.size() - 1);
        res += sum[lim];
    }
    cout << res << endl;
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AOJ 3182 Umg Kart (HUPC 2020 day3-K)

問題概要

制約

考えたこと

距離 1 あたり t 秒で走る選手のゴール時間の確率分布

部分和 DP を求める問題を FFT へ

確率分布から答えへ

コード