Rustコトハジメ

プログラミング言語Rustと競プロに関する情報をお届けします。

ビットDP良問「ぬいぐるみの整理」

D - ぬいぐるみの整理 (Plush Toys)

今日の勉強メモです。ビットDPはたまに出てくるのと、わりと苦手意識があるので、良問っぽいのを実装訓練も兼ねて解いておいた。累積和とビットDPの組み合わせ。

問題

M(20)種類の数字が適当にN(105)個並べてある。このうちいくつかを一旦取り除いたのち、並び替えて空き地に戻す。この時、同じ種類の数字が連続している(3333111222みたいな感じ)ようにしたい。取り除く最小の個数はいくつか?

例: 1222121なら、1と2を1つずつ取り除いて2222111にするのが最善。

考え方

O(M! N)とかいう超絶クソ解法は思いついたが、当然間に合わない。この考えは、種類の並び順とそれぞれの種類の個数がわかれば、数字の並び順はわかるから、差分で取り除く個数がわかる。全ループすれば最小はわかる。というもの。これを基本として考える。

ポイント1 まず、仮に並び順がわかったとする。すると、各種類の連続列と、初期列との差分は、その連続列の範囲にその種類があった個数と考えることが出来る。これは、累積和を計算しておくことでO(1)で計算可能となる。

ポイント2 次に、種類をDPで拡張していくという方針にする。つまり、種類の集合Sを使っている時に、何か他の種類をくっつけていくことを考える。この時、その種類について取り除く個数(ソースコードでいうところのk_removal)は上の累積和から計算出来る。

実装

fn solve() {
    input! {
        N: usize, M: usize,
        X: [usize; N],
    }

    let mut kind_count = vec![0; M];
    let mut kind_bit = vec![vec![0; N+1]; M];
    let mut kind_count_acc = vec![vec![0; N+1]; M];
    // O(NM)
    for k in 0..M {
        for i in 0..N {
            if X[i] == k+1 {
                kind_count[k] += 1;
                kind_bit[k][i] = 1;
            }
        }
    }
    // O(NM)
    for k in 0..M {
        let mut acc = 0;
        for i in 0..N {
            if kind_bit[k][i] == 1 {
                acc += 1;
            }
            kind_count_acc[k][i+1] = acc;
        }
        // dbg!(&kind_count_acc);
    }
    let mut dp = vec![1<<40; 1<<M];
    dp[0] = 0;
    // O(2^M)
    for s in 0..(1<<M) {
        let mut s_tail = 0;
        // O(M)
        for k in 0..M {
            if s & (1<<k) > 0 {
                s_tail += kind_count[k];
            }
        }
        // O(M)
        for k in 0..M {
            if s & (1<<k) > 0 {
                continue;
            }
            let k_keep = kind_count_acc[k][s_tail + kind_count[k]] - kind_count_acc[k][s_tail];
            let k_removal = kind_count[k] - k_keep;
            dp[s | (1<<k)] = min(
                dp[s | (1<<k)],
                dp[s] + k_removal
            );
            // dbg!(s,k,s|(1<<k),k_removal,&dp);
        }
    }
    // dbg!(&dp);
    println!("{}", dp[(1<<M)-1]);
}

考察

Mが小さいのでビットDPぽさ(あるいは半分全列挙)はある。そう考えると、「すでに〜した」と巡回セールスマンと同じように考えて、「すでに使った種類」についてビットDPをすれば良いということはわかるかも知れないが、そこから累積和を使ってDPを出来るようになる流れが美しく感じた。あまりに美しいので、ビットDP+累積和って典型なのかと疑っている。

巡回セールスマンについてはこちら。

巡回セールスマン問題の考え方 - Rustコトハジメ

累積和は、[l,r)についての累積和が、acc[r]-acc[l]で求まるようにするには、インデックスを一つずらす必要がある。これが、

            if kind_bit[k][i] == 1 {
                acc += 1;
            }
            kind_count_acc[k][i+1] = acc;

の意味。この問題の場合、kind_countという長さごとに伸ばしていくため、[l,r)で問い合わせるのが実装として無理がなく、上のように累積和を工夫するのが良かった。