ARC073 Many Moves - tinumu's reminder

atcoder.jp
ACPCVC20200120のCで出ました
このセグ木の使い方が少し教育的と感じたため、メモとして記述してみます。

問題文

$N$ 個のマスがある。コマを２つ持っていて、コマはそれぞれマス $A, B$ にある。
マス $x_i$ にどちらかのコマを移動させるというクエリを $Q$ 個処理する。
クエリを処理した後の、コマの移動距離の最小を求める。

$1 \leq N, Q \leq 2 \times 10^5$
$1 \leq A, B \leq N$
$1 \leq x_i \leq N$

解法

1

それぞれのコマが今どのマスにいるかという情報がわかれば、この先の移動の仕方は同じになることがわかるので、まず以下のDPが作れる。
$dp[i][a][b]$ : $i$ 個のクエリを処理し終え、現在のコマは $a, b$ に存在する。
現在は $O(QN^2)$ となっている。まだ遅い。

2

次に、どちらかのコマの座標は $x_i$ になっていることに着目すると、実はどちらかのコマの情報がいらない。これで変数が2つになる。
$dp[i][a]$ : $i$ 個のクエリを処理し終え、現在のコマは $a, x_i$ に存在する。
遷移も後に説明するが、ちゃんと $O(1)$ で出来るため、 $O(QN)$ になった。まだ。

3

実はここからの高速化は、DPの遷移自体を小さくするわけではない。一気に $N$ 個遷移することで高速化する。

ところで、先ほどのDPの遷移はどうなっていたのか。それは次のようになっている。

$dp[i][a]=dp[i-1][a]+|x_i-x_{i-1}|(a \neq x_{i-1})$
$dp[i][x_{i-1}]=min(dp[i-1][x_{i-1}]+|x_i-x_{i-1}|,min_{j=1,2,...,n}(dp[i-1][j]+|x_i-j|))$

一見複雑そうに見えるが、実はそうでもない。
$x_{i-1}$ を $x_i$ に動かすか、どこかにある $j$ を $x_i$ に動かすかの2択しか無い。
ちゃんと睨めば漸化式が出てくるはず。

ここで、 $x_{i-1}$ を $x_i$ に動かす動作の遷移を見てみると次のようになっている。

…

□□□□□□□□□□□□□□□□□□□□□□□

↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓↓

□□□□□□□□□□□□□□□□□□□□□□□

…

$i-1$ から $i$ に遷移しているだけ。しかも遷移にかかるコストは $|x_i-x_{i-1}|$ で一定となっており配列全体に $|x_i-x_{i-1}|$ を足すという処理で $N$ 個の遷移を実現できる。
配列全体に同じものを足す遷移は高速にできるので、適宜工夫すると良い。この操作自体は $O(1)$ で行える(実際はSegmentTreeの影響で $O(logN)$ 掛かりそう)

4

次に、 $min_{j=1,2,...,n}(dp[i-1][j]+|x_i-j|)$ の高速な計算方法を説明する。

この式は絶対値を外すことで、扱いやすい形になる。

$min_{j=1,2,...,n}(dp[i-1][j]+|x_i-j|)=$
$min(min_{j=1,2,...,x_i}(dp[i-1][j]-j)+x_i, min_{j=x_i+1,x_i+2...,n}(dp[i-1][j]+j)-x_i)$

このとき、 $dp[i-1][j]-j$ と $dp[i-1][j]+j$ の配列を持っておけば、区間のminを求める操作でこの処理を実現できることがわかる。
これは予めdpのマス $j$ に対して $j$ や $-j$ を足しておくことで作ることが出来る。
区間のminを求める操作はSegment Tree を使うと $O(logN)$ で動作するので高速に動作することが出来る。(2個セグ木をつかう)

これで、全体で $O(QlogN)$ となり、間に合わせることが出来る。

このDPをインラインDPというらしい。
topcoder.g.hatena.ne.jp
とても素晴らしいサイトでした。

ソースコード

int main() {
    ll N, Q, A, B;
    ll X[200005];
    cin >> N >> Q >> A >> B;
    SegmentTree<ll> segl(N + 3, LLINF, [](ll a, ll b) { return (min(a, b)); });
    SegmentTree<ll> segr(N + 3, LLINF, [](ll a, ll b) { return (min(a, b)); });
 
    ll sum = 0;
 
    X[0] = A;
    segl.update(B, -B);
    segr.update(B, B);
    for (int i = 1; i <= Q; i++) {
        cin >> X[i];
        ll diff = llabs(X[i] - X[i - 1]);
        sum += diff;
        ll val = min(segl.query(1, X[i] + 1) + X[i], segr.query(X[i] + 1, N + 1) - X[i]);
        segl.update(X[i - 1], val - X[i - 1] - diff);
        segr.update(X[i - 1], val + X[i - 1] - diff);
    }
    ll minv = 0;
    for (int i = 1; i <= N; i++) {
        minv = min(minv, segl.query(i, i + 1) + i);
    }
    cout << minv + sum << endl;
 
    return (0);
}

感想

2次元DPの遷移を一気に $N$ 回やるっていう方法が知れてよかった。
DP、使えるなあ。
あと、ソースコードのフォントサイズ、どうにかならないかな。