🐰

Rustで動的計画法：Grouping

2023/05/31に公開

Rust

AtCoder

動的計画法

tech

はじめに

動的計画法を実装してみて、Rustの勉強をやってみる。問題としてはEducational DP Contestという動的計画法の練習を目的としたコンテストのものを使用。

https://atcoder.jp/contests/dp/

AからZまで問題が設定されているが、今回はUのGroupingを解いてみる。うさぎを分類して一番高い組み合わせを見つける問題。

利用したライブラリ

標準入力からデータを読み取る部分はproconioを利用。

総当たりでの実装

うさぎの数は $1\leq N\leq 16$ となっているので、あきらかにビット演算で解くのだと思うが、
まずは、普通に解いてみる。基本的に総当たりですべての組み合わせを調べていく。

grouping.rs

use proconio::input;
use itertools::Itertools;
use std::collections::HashMap;

fn powerset(s: &Vec<usize>, i: usize, j:usize) -> impl Iterator<Item=Vec<usize>> + '_ {
    (i..=j).map(|c| s.clone().into_iter().combinations(c)).flatten()
}

fn difference(s: &Vec<usize>, t: &Vec<usize>) -> Vec<usize> {
    s.clone().into_iter().filter(|x| !t.contains(x)).collect()
}

fn main(){
    input! {
        n : usize,
        a : [[i64; n]; n],
    }

    let rabbit = Vec::from_iter(0..n);
    let mut dp =  HashMap::<Vec<usize>, i64>::new();

    for i in &rabbit {
        dp.insert(vec![*i], 0i64);
    }

    for i in powerset(&rabbit, 2, n) {
        let mut value = 0i64;
        for e in i.iter().combinations(2) {
            let (r1, r2) = (*e[0], *e[1]);
            value += a[r1][r2];
        }
        for e in powerset(&i, 1, i.len() - 1) {
            let f: Vec<usize> = difference(&i, &e);
            value = std::cmp::max(value, dp.get(&e).unwrap() + dp.get(&f).unwrap());
        }
        dp.insert(i, value);
    }
    let ans = dp.get(&rabbit).unwrap();
    println!("{:?}", ans);
}

解き方

あるウサギのセットiに対して、最高得点dp[i]を導出することを考える。まずは、すべてのウサギを1つのgroupに入れた時の点数を計算する。 $i$ から2つを選んだすべての組み合わせについて、点数を足し合わせる事で求めている。

for e in i.iter().combinations(2) {
    let (r1, r2) = (*e[0], *e[1]);
    value += a[r1][r2];
}

次にセットiをeとf2つに分けてそれぞれの最大値の和から、そこで切った時の最大値が分かる。この切り方は、 $i$ のすべての部分集合、つまり「冪集合(Power Set)」を求めてそれでループをかける。すべての切り方の中で最も点数の高い切り方の点数がdp[i]となる。

eとfはiの部分集合なので、要素数はiより少ない。このため要素数の少ない方からループを回せば、答えを計算できる。

for e in powerset(&i, 1, i.len() - 1) {
        let f: Vec<usize> = difference(&i, &e);
        value = std::cmp::max(value, dp.get(&e).unwrap() + dp.get(&f).unwrap());
}

冪集合の求め方はStack Overflowにでているが、今回は要素数を変えてcombinationを呼び出す方法をとった。こうすると要素数が少ない順に並んだ行列が取得できる。

(i..=j).map(|c| s.clone().into_iter().combinations(c)).flatten()

イメージとしては以下の感じ

for c in 1..=s.len() {
    for i in s.iter().combinations(k) {

iとjを引数にすることで、要素数が1の部分集合や入力の行列そのものを除いている。

Iteratorを返す関数

powersetはそのままループに使うのでIteratorを返す事にした。Iteratorは型ではなくTraitである。今回の場合以下のような型になっている（みたい）。

Flatten<Map<RangeInclusive<usize>>>

これを書いていると大変なので、Vec<usize>のIteratorが実装された型を返す関数ということを表すために impl Iterator<Item=Vec<usize>>という書き方をする。

fn powerset(s: &Vec<usize>, i: usize, j:usize) -> impl Iterator<Item=Vec<usize>> + '_ {
　　　　(i..=j).map(|c| s.clone().into_iter().combinations(c)).flatten()
}

後半でIteratorの実装を行うが、Iteratorはすぐに評価されない。返り値に参照が存在しないが、内部的にはsへの参照を持っているので、それを示すために匿名ライフタイム `_を入れる必要がある。

コンパイラが注意できるので、どちらかと言えば「読む人が間違えないように」のような気はする。明示的に書くとすると以下のような感じだと思う。

fn powerset<'a>(s: &'a Vec<usize>, i: usize, j:usize) -> impl Iterator<Item=Vec<usize>> + 'a {

ビット演算を使った実装

処理としては上記のものと同じであるが、こっちが本命であり、制限時間内に終わるのはこちら。ただ、ローカル環境での試験でも結構ギリギリなので、まだ無駄が多いのかもしれない。

grouping_bit.rs

use proconio::input;
use itertools::Itertools;

struct Powerset {
    set: usize,
    next: usize,
}

impl Powerset {
    fn new(s: usize) -> Self {
        Powerset { set: s, next: s }
    }
}

impl Iterator for Powerset {
    type Item = usize;

    fn next(&mut self) -> Option<Self::Item> {
        if self.next == 0 {
            None
        } else {
            let current = self.next;
            self.next = (current - 1) & self.set;
            Some(current)
        }
    }
}


fn main() {
    input! {
        n : usize,
        a : [[i64; n]; n],
    }

    let rabbit = (1usize << n) - 1;

    let mut comp = vec![0i64; rabbit + 1];
    for c in (0..n).combinations(2) {
        let (x, y) = (c[0], c[1]);
        comp[1usize << x | 1usize << y] = a[x][y];
    }    

    let mut dp = vec![0i64; rabbit + 1];
    for i in Powerset::new(rabbit).sorted_by(|a, b| a.count_ones().cmp(&b.count_ones())) {
        if i.count_ones() > 1 {
            let mut value = 0i64;
            for j in Powerset::new(i).filter(|x| x.count_ones() == 2) {
                value += comp[j];
            }
            for j in Powerset::new(i) {
                if j != i {
                    let k = j ^ i;
                    value = std::cmp::max(value, dp[j] + dp[k]);
                }
            }
            dp[i] = value;
        }
    }
    let ans = dp[rabbit];
    println!("{:?}", ans);
}

Iteratorの実装

こちらではPowersetはIteratorとして実装した。

Iteratorの実装は難しくなく、fn next(&mut self) -> Option<Self::Item>を実装していれば良い。nextが呼ばれる度に次の値をSome(v)の形で返して、要素がなくなったらNoneを返すと良い。for文で使った場合は、ループ毎に呼び出されるが、今回の場合要素数が少ない方からループさせるためにsort_byしているので、最初に全要素を計算してしまっている。