Rustコトハジメ

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

【Educational DP-U Grouping】ビットDPと足し合わせ高速化

U - Grouping

問題

N個のノードがあり、各ノードには得点Aijが決められている。

ノードをいくつかのグループに分割する。得点の最大値を求めよ。

制約: N=16, |Aij|=109

方針

E - Get Everything

の強化版です。というのも事前計算をする必要があるのと、Nが少し大きいため、愚直に二重ループすると間に合わないからです。

基本的には、グループをマスクで表して、それぞれについて得点を計算(O(N2 2N)なので間に合う)し、合計が111...111になるようなグループのうち、得点最大をDPで探します。

ただし、216 = 6*104なので、これを二重ループで回すとTLEします。

そこで各マスクについて、とりうる遷移を絞ることにします。こうすると、

 \sum nCk 2k

となるはずで、正確にいくつになるかは知りませんが、間に合います。(これがどのくらいで抑えられるのか知ってる人は教えてください。@GodAimAkira)直感的には、2nの時にnCk=1なので「だいぶ小さくなった」ということは言えそうですが。

とりあえず計算してみると、以下のような倍率であることはわかりました。nが16の時は2桁くらい速いということが言えそうです。今回でいうと元が9乗なので7乗になってぎりぎり間に合うという感じでしょうか。

n 倍率
1 1
2 1
3 2
4 3
5 4
6 5
7 7
8 9
9 13
10 17
11 23
12 31
13 42
14 56
15 74
16 99
17 133
18 177
19 236
20 315

実装

fn solve() {
    input!{
        N:usize,
        a:[[i64;N];N],
    }
    // cost[bit]
    let mut cost = vec![0; 1<<N];
    for bit in 0..1<<N {
        let mut hs = HashSet::new();
        for i in 0..N {
            if bit&(1<<i) > 0 {
                hs.insert(i);
            }
        }
        let mut c: i64 = 0;
        for x in &hs {
            for y in &hs {
                if x>y { continue; }
                c += a[*x][*y];
            }
        }
        cost[bit]=c;
    }

    let mut dp = vec![0;1<<N];
    for bit in 0..1<<N {
        let mut x = bit;
        loop {
            if x == 0 { break; }
            dp[bit]=max(dp[bit], dp[bit^x] + cost[x]);
            x = (x-1)&bit;
        }
    }
    println!("{}",dp[(1<<N)-1]);
}

考察

最後に使ってる謎のビット操作は、元のマスクが持ってるビットを使ったビットマスクの全通りを計算しています。ライブラリ化しておきました。

impl Iterator for AllMasks {
    type Item = u64;
    fn next(&mut self) -> Option<Self::Item> {
        let old = self.smask;
        if old == 0 {
            return None
        }
        self.smask = (self.smask-1) & self.mask;
        return Some(old)
    }
}

フル実装はこちら。

GitHub - akiradeveloper/rust-comp-snippets