JOIG 春合宿 2023 day2-2 White Light (2D, 難易度 8)

セグメント木を用いた DP 高速化！

• 問題へのリンク
• editorials

問題概要

'R' と 'G' と 'B' のみからなる長さ の文字列 が与えられる。以下の操作を繰り返し行うことで、"RGB" を繰り返す文字列となるようにしたい。

(操作)
連続する 個以下の文字を消す

目的を達成するために必要な操作回数の最小値を求めよ。

制約

• 

考えたこと

この問題のように「区間に対して操作する」というタイプの問題では、DP がうまくいくことが非常に多い。まずは計算量を気にせずに DP をしてみよう。

次の配列 dp を定義する。

---

dp[c][i] ← 文字列 の先頭から 文字分のみを考える。それらから問題文で定義された操作を実行後に、最後の文字が次のようになっている状態にするために必要な操作の最小回数

• c = 0 のときは 'R'
• c = 1 のときは 'G'
• c = 2 のときは 'B'

---

このとき、テクニカルだが、初期条件は dp[2][0] = 0 とすればよい ('B' で終わっているから次は 'R' から始まるというように考える)。

それでは、 に対する dp[0][k], dp[1][k], dp[2][k] の値を用いて、dp[0][i+1], dp[1][i+1], dp[2][i+1] を更新する式を考えよう。

文字 S[i] を最後に使う場合

文字 S[i] に対応する数値を c とする。このとき、次のように更新できる。

• dp[c][i+1] = min(dp[c][i+1], dp[(c+2)%3][i])

ここで、(c+2)%3 は、数値 c に対応する文字の前の文字に対応する数値を表している。

文字 S[i] を使わない場合

この場合は、S[i] を含む何文字かを連続して消す操作を最後に行うこととなる。その文字数を () とすると、次の更新式が立てられる。

• dp[0][i+1] = min(dp[0][i+1], dp[0][i+1-k] + 1)
• dp[1][i+1] = min(dp[1][i+1], dp[1][i+1-k] + 1)
• dp[2][i+1] = min(dp[2][i+1], dp[2][i+1-k] + 1)

計算量

この時点で、DP の更新に要する計算量は となる。ここまでで部分点 62 点がとれる。

セグメント木で DP 高速化

それでは、この DP を高速化しよう。

• dp[c][i+1] = min(dp[c][i+1], dp[(c+2)%3][i])

の方は現状のままでも問題ない。問題なのは、

• dp[0][i+1] = min(dp[0][i+1], dp[0][i+1-k] + 1)
• dp[1][i+1] = min(dp[1][i+1], dp[1][i+1-k] + 1)
• dp[2][i+1] = min(dp[2][i+1], dp[2][i+1-k] + 1)

の方である。この更新を について実施しているから の計算量を要するのである。ここで、次のような関数を考えることとしよう。

• dp[0][k]
• dp[1][k]
• dp[2][k]

これを用いると、上記の計算式は次のように書ける。

---

• dp[0][i+1] = min(dp[0][i+1], f(i+1-K, i+1) + 1)
• dp[1][i+1] = min(dp[1][i+1], g(i+1-K, i+1) + 1)
• dp[2][i+1] = min(dp[2][i+1], h(i+1-K, i+1) + 1)

---

もし、 を指定したときに の値を高速に計算できる方法があれば、上記の DP 更新も高速にできる。その方法としてセグメント木がある。セグメント木を用いれば、配列の区間の最小値を求める操作を の計算量で実現できる。

セグメント木については、まずは次の問題で練習しよう。

drken1215.hatenablog.com

この問題に戻ると、上述の DP の更新をセグメント木上で実現することで、全体の計算量は となる。

コード

#include <bits/stdc++.h>
using namespace std;
const int INF = 1<<29;

// 区間の最小値を返すセグメント木 (最低限の実装)
struct SegmentTree {
    int log, offset;
    vector<int> dat;
    
    SegmentTree(int N) {
        log = 0, offset = 1;
        while (offset < N) ++log, offset <<= 1;
        dat.assign(offset*2, INF);
    }
    
    void set(int i, int val) {
        i += offset;
        dat[i] = val;
        while (i >>= 1) dat[i] = min(dat[i*2], dat[i*2 + 1]);
    }
    
    int prod(int l, int r) {
        int val_left = INF, val_right = INF;
        l += offset, r += offset;
        for (; l < r; l >>= 1, r >>= 1) {
            if (l & 1) val_left = min(val_left, dat[l++]);
            if (r & 1) val_right = min(dat[--r], val_right);
        }
        return min(val_left, val_right);
    }
    
    int operator [] (int i) {
        return prod(i, i+1);
    }
};

int main() {
    int N, K;
    string S;
    cin >> N >> K >> S;
    
    vector<SegmentTree> dp(3, SegmentTree(N + 1));
    dp[2].set(0, 0);
    for (int i = 0; i < N; ++i) {
        int c = 0;
        if (S[i] == 'G') c = 1;
        else if (S[i] == 'B') c = 2;
        
        // 文字 S[i-1] を採用する場合
        dp[c].set(i+1, min(dp[c][i+1], dp[(c+2)%3][i]));
        
        // 最初の i+1-j (j = 1, 2, ..., K) 文字分の答え + 1
        for (int j = 0; j < 3; ++j) {
            dp[j].set(i+1, min(dp[j][i+1], dp[j].prod(max(i+1-K, 0), i+1) + 1));
        }
    }
    cout << dp[2][N] << endl;
}