AtCoder Beginner Contest 158 F - Removing Robots

問題: https://atcoder.jp/contests/abc156/tasks/abc158_f

問題概要

数直線上に $N$ 台のロボットが置かれている。$i$ 番目のロボットの初期位置は $X_i$ で，このロボットは起動させると $X_i+D_i$ まで移動する。移動中，$i$ 番目のロボットが $X_i \le X \lt X_i+D_i$ を満たす $X$ にいるとき，$ X_j = X$ を満たす他のロボット (このロボットを $j$ 番目のロボットとする) が存在すれば，$i$ 番目のロボットと $j$ 番目のロボットは接触したと考える。このとき，接触された $j$ 番目のロボットは起動し移動し始める。

高橋くんは次の操作を何回でも行うことができる。

数直線上のロボットから1台のロボットを選び，起動させる。

移動中のロボットがいるときは高橋くんはロボットを起動させることはできず，移動を終えたロボットは数直線上から取り除かれるとする。

高橋くんが操作を終えた後に数直線上に残っているロボットの組み合わせとして考えられるものの通り数を $998244353$ で割った余りを答えよ。

解法概要

数直線上に残っているロボットの組み合わせを数えるのは取り除かれるロボットの組み合わせを数えるのと等しいので，以下，取り除かれるロボットの組み合わせを考える。

また，ロボットは $X_i$の値によってソート済みとする。

仮に，$0 \leq i \lt N$ のすべての $i$ について，$i$ 番目のロボットを起動させたときにロボット同士の接触で起動されるロボットが $j (\geq i)$ 番目までのロボットであるとわかっているとすると，次のような DP で求めるべき通り数を計算することができる。 ($i$ 番目のロボットとの接触とは限らない点に注意。$i$ 番目のロボットと接触して動き出した $j$ 番目のロボットが，$i$ 番目のロボットとは接触しない別の $k$ 番目のロボットと接触して $k$ 番目のロボットが起動する，といったことがありえる)

$\mathit{dp}_i := i$ 番目以降のロボットの残りかたの通り数

$\mathit{dp}_{n+1} = 1$ (いずれのロボットも起動させない場合に対応する)

$\mathit{dp}_i = \mathit{dp}_{i+1} + \sum_{j'=j}^{n+1} \mathit{dp}_{j'}$ (第 1 項は $i$ 番目のロボットを残す場合の $i+1$ 番目以降のロボットの残りかたの通り数，第 2 項は $i$ 番目のロボットを起動させる場合の通り数)

上式の第2項について，$i+1$ 番目のロボットから $j-1$ 番目のロボットについては，$i$ 番目のロボットを起動させると衝突によって強制的に起動されてしまうので足しこむ必要はない。

さて，肝心の「$i$ 番目のロボットを起動させたときにロボット同士の接触で起動されるロボットが $j (\geq i)$ 番目までのロボットである」の $j$ を求めるにはどのようにすればよいだろうか。

すぐに思いつくのが，「$X_j \geq X_i + D_i$ となる最初の index $j$ を二分探索で求め，そこから 1 を引いたもの」であるがこれは実は間違っている。例えば，

のような場合，1 番目のロボットに影響を受けるのは 2 番目のロボットと 3 番目のロボットの両方であるが，上のやり方だと 2 番目のみということになってしまう。

ではどうすればよいのかというと，これも DP っぽく計算できる。ソースコードに合わせて， $j$ ではなく $j+1$ にあたるもの (= $i$ 番目のロボットが起動されても影響を受けないロボットのうち座標が最も小さいものの index) を $\mathit{ti}_i$ に格納することを考える。

$ k = X_j \geq X_i + D_i$ となる最初の index とすると，

$\mathit{ti}_{n+1} = n+1$

$\mathit{ti}_i = \max_{i \lt i' \le tmp_k} \mathit{ti}_{i'}$

となる。第2式は，$i \lt i'$ となる$i'$ について $\mathit{ti}_{i'}$ が正しく求められているとすると，$i$ 番目のロボットと直接衝突するロボットについて，どこまで影響が及ぶのかの最大範囲を見ればよいということに対応する。最大値をとるときに毎回愚直にループを回していると処理全体で時間計算量が $O(N^2)$ に及んでしまうので，セグ木などを使って $O(n\log n)$ に抑えよう。

あとは求めた $\mathit{ti}_i$ を使って最初の DP を計算すればよい。

ソースコード

#include <iostream>
#include <vector>
#include <algorithm>
#include <functional>

using namespace std;

using ll =  long long;
using Pll = pair<ll, ll>;

constexpr ll MOD = 998244353;

template<class T> class SegmentTree {
    public:
    typedef function<T(T, T)> F;
    size_t n;
    vector<T> data;
    T def;  // 単位元
    T initial_value;    // 初期値
    F update_func;  // 更新用関数
    F find_func;  //クエリ用関数

    SegmentTree(size_t _n, T def, T initial_value, F update_func, F find_func)
        : def(def), initial_value(initial_value), update_func(update_func), find_func(find_func) {
        n = 1;
        while(n < _n) n <<= 1;
        data = vector<T>(2*n-1, initial_value);
    }

    void update(size_t i, T x) {
        i += n-1;
        data[i] = update_func(data[i], x);
        while(i) {
            i = (i-1) / 2;
            data[i] = find_func(data[2*i+1], data[2*i+2]);
        }
    }

    T find(size_t s, size_t t, size_t k, size_t kl, size_t kr) {
        if(kr <= s || t <= kl) return initial_value;
        if(s <= kl && kr <= t) return data[k];
        size_t kc = (kl+kr) / 2;
        T vl = find(s, t, 2*k+1, kl, kc);
        T vr = find(s, t, 2*k+2, kc, kr);
        return find_func(vl, vr);
    }

    T find(size_t s, size_t t) {
        return find(s, t, 0, 0, n);
    }
};

int main() {
    ios::sync_with_stdio(false); cin.tie(nullptr);
    int n;
    cin >> n;
    vector<Pll> robots(n);
    for(int i=0;i<n;++i) cin >> robots[i].first >> robots[i].second;
    sort(robots.begin(), robots.end());

    // 初めて影響が及ばなくなる idx
    SegmentTree<int> ti(
        n+1, 0, 0, 
        [](ll a, ll b) { return b; },
        [](ll a, ll b) { return max(a, b); }
    );
    ti.find(n, n);
    for(int si=n-1;si>=0;--si) {
        int tmp_ti = int(lower_bound(robots.begin()+si, robots.end(), Pll(robots[si].first + robots[si].second, -1LL)) - robots.begin());
        if(tmp_ti - 1 > si) tmp_ti = ti.find(si+1, tmp_ti);
        ti.update(si, tmp_ti);
    }

    // si 以降を見たとした場合の通り数
    vector<ll> dp(n+1, 0);
    dp[n] = 1;
    for(int si=n-1;si>=0;--si) {
        dp[si] = (dp[ti.find(si, si+1)] + dp[si+1]) % MOD;
    }

    cout << dp[0] << endl;
}

所感

セグ木使った DP の高速化すき

好きなタイプ以外の問題の解説もちゃんと書いた方が勉強になるんだけど，ついこういう問題の解説を優先してしまうね