🎰

Rustで動的計画法：Permutation

2023/05/26に公開

Rust

AtCoder

動的計画法

tech

はじめに

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

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

AからZまで問題が設定されているが、今回はTのPermutationを解いてみる。順列を文字列 $s$ に従って並び替える問題。解き方が難しいので、行列の操作を学びながら総当たりからやってみる。

利用したライブラリ

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

総当たりでの実装

まずは総当たりでやってみる。深さ優先探索で順列remainから条件にあう数値を取り出しつつ、進めていく。

permutation_1.rs

use proconio::input;
use proconio::marker::Chars;

const MOD:usize = 10usize.pow(9) + 7;

fn check(s: char, a: usize, b: usize) -> bool {
    if s == '>' { a > b } else { a < b }
}

fn solve(remain: &[usize], str: &[char], str_pos: usize, priv_value: usize) -> usize {
    let mut count = 0;
    for (index, value) in remain.iter().enumerate() {
        if check(str[str_pos].clone(), priv_value, *value) {
            if remain.len() > 1 {
                let mut r = remain.to_vec();
                r.swap_remove(index);
                count = (count + solve(&r, str, str_pos + 1, *value)) % MOD;
            } else {
                count = (count + 1) % MOD;
            }
        }
    }
    count
}


fn main(){
    input! {
        n: usize,
        s: Chars,
    }

    let array = Vec::from_iter(0..n);
    let mut ans = 0;
    for i in 0..n {
        let mut remain = array.clone();
        remain.swap_remove(i);
        ans = (ans + solve(&remain, &s, 0, i)) % MOD;
    }
    println!("{}", ans);
}

ベクタと行列とスライス

RustではベクタVec<T>と行列[T; N]がある。ベクタは可変長だが、行列はコンパイル時に長さが決まっていないといけない固定長。「コンパイル時に長さが決まっていないといけない」というのがネックで使い所が難しい気がする。

これとは別にスライス&[T]というものがありこちらは、データコレクションの一部の要素を示す。関数の引数などに使われることが多いイメージだがが、これが分かりにくい...というか分かったつもりでいても実際使うとエラーがでることが多い。

今回は、関数内で変更しない（immutable）ベクタを渡す場合は、スライスで渡すとベクタでも行列でも関数を呼び出すことができるという観点で使っている。つまりここは

fn solve(remain: &Vec<usize>, str: &Vec<char>, str_pos: usize, priv_value: usize) -> usize {

と、ベクタの参照を渡してもいいし、remainは関数呼び出しのために作っているので

fn solve(remain: Vec<usize>, str: &Vec<char>, str_pos: usize, priv_value: usize) -> usize {

と、所有権ごと渡しても構わない。ただ、ベクタを関数の引数にしてしまうとコピーが発生しそうに見えるので参照渡し、さらには好みかもしれないが、スライス渡しの方が見通しはいいのかもしれない。とはいいつつ🐍LCSではベクタの参照渡しで書いているが。

スライスはcloneできないので、スライスからベクタを作りたいときはto_vec()を使う。ここがちょっと直感的な書き方でないのが少し気になる。

let mut r = remain.to_vec();

removeとswap_remove

ベクタから特定の位置の要素を削除するときはremoveを使う。

r.remove(index);

ただ、このときindexより後ろの要素を前にずらす必要があるので処理としては $O(n)$ かかってしまう。もし、順番を気にしないのであればswap_removeを使うと良い。削除したindexの位置に、ベクタの最後尾の要素を持ってくるので、 $O(1)$ で処理が終わる。

2つのベクタを渡す

上のコードを見ると、remainの中から条件を見たす要素を見つけて、それをpriv_valueとし、次の条件、つまりpriv_valueより大きいのか小さいのかに使っている。ただし、数列は降順にならんでいるため、priv_valueを渡す代わりに、それよりも大きい残りの要素、小さい残りの要素を渡す事は容易である。そこで、以下のように修正する。

permutation_2.rs

use proconio::input;
use proconio::marker::Chars;

const MOD:usize = 10usize.pow(9) + 7;

fn solve(left: &[usize], right: &[usize], str: &[char], str_pos: usize) -> usize {
    if str_pos == str.len() {
        1
    } else {
        let mut count = 0;
        if str[str_pos] == '>' {
            for i in 0..left.len() {
                let mut l = left.to_vec();
                let mut r = right.to_vec();
                let mut c = l.split_off(i + 1);
                l.pop();
                c.append(&mut r);
                count = (count + solve(&l, &c, str, str_pos + 1)) % MOD;
            }
        } else {
            for i in 0..right.len() {
                let mut l = left.to_vec();
                let mut c = right.to_vec();
                let r = c.split_off(i + 1);
                c.pop();
                l.append(&mut c);
                count = (count + solve(&l, &r, str, str_pos + 1)) % MOD;
            }
        }
        count
    }
}


fn main(){
    input! {
        n: usize,
        s: Chars,
    }

    let array = Vec::from_iter(0..n);
    let mut ans = 0;
    for i in 1..n {
        let mut left = array.clone();
        let right = left.split_off(i);
        left.pop();
        ans = (ans + solve(&left, &right, &s, 0)) % MOD;
    }
    println!("{}", ans);
}

ベクタの分割と結合

ベクタを分割するにはsplit_off()を使う。i番目以上の要素を切り出してベクタにして返す。

let mut a = vec![1, 2, 3, 4, 5];
let b = a.split_off(3);
println!("{:?} {:?}", a, b);
// [1, 2, 3] [4, 5]

結合するときはappendを使う。append元のベクタの中身は空っぽになる。

let mut a = vec![1, 2, 3];
let mut b = vec![4, 5];
a.append(&mut b);
println!("{:?} {:?}", a, b);
// [1, 2, 3, 4, 5] []

前の数値より小さい数値の数で考える

permutation_2.rsをよく見ると、rightやleftの中に何が入っているかはあまり関係なく、前に設置した要素より小さい要素が何個残っているのか、大きい要素が何個残っているかだけしか見ていない。つまり、それぞれの行列の要素数だけで処理ができる。具体的にsolve()を書き直すと以下のようになる。

fn solve

fn solve(left: usize, right: usize, str: &[char], str_pos: usize) -> usize {
    if str_pos == str.len() {
        1
    } else {
        let mut count = 0;
        if str[str_pos] == '>' {
            for i in 0..left {
                count = (count + solve(left - i - 1, right + i, str, str_pos + 1)) % MOD;
            }
        } else {
            for i in 0..right {
                count = (count + solve(left + i, right - i - 1, str, str_pos + 1)) % MOD;
            }
        }
        count
    }
}

動的計画法で解く

これまでは、総当たりで解いており、 $O(N^N)$ になる。このため、 $N$ が少し大きくなると処理に時間がかかる。そこで最後に行った「前の数値より小さい数値の数」をベースとして動的計画法を組んで見る。

🍬Candiesでやったように累積和を使うことを忘れなければ容易に実装できる。計算量としては $O(N^2)$ となる。

permutation_3.rs

use proconio::input;
use proconio::marker::Chars;

const MOD:usize = 10usize.pow(9) + 7;

fn main(){
    input! {
        n: usize,
        s: Chars,
    }

    let mut dp = vec![1usize; n];
    for c in &s {
        let l = dp.len() ;
        let mut next_dp = vec![0usize; l - 1];
        if *c == '>' {
            for j in 1..l {
                next_dp[0] = (next_dp[0] + dp[j]) % MOD;
                if j < l - 1 {
                   next_dp[j] = (next_dp[j] + MOD - dp[j]) % MOD;
                }
            }
        } else {
            for j in 0..(l - 1) {
                next_dp[j] = (next_dp[j] + dp[j]) % MOD;
            }
        }
        for j in 1..(l - 1) {
            next_dp[j] = (next_dp[j] + next_dp[j - 1]) % MOD;
        }
        dp = next_dp;
    }
    let ans = dp[0];
    println!("{}", ans);
}

過去の記事

Rustで動的計画法の実装: