🌻

Rustで動的計画法：Flowers

2023/05/12に公開

Rust

AtCoder

動的計画法

tech

はじめに

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

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

AからZまで問題が設定されているが、今回はQのFlowersを解いてみる。 $n$ 本の横一列に並んでいる花を抜いて高さが単調増加になるようにするというどういう状況がよくわからない問題である。

利用したライブラリ

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

素直に解いてみたコード

まずは、素直に $i$ 本目が一番左側の花だった時の最大の美しさをdp[i]として解いてみる。分かりやすいが、計算量は $O(n^2)$ 。 $1 \leq N \leq 2\times 10^5$ なので、所定時間内に終わらない。

flowers_loop.rs

use proconio::input;
use std::cmp::max;

fn main(){
    input! {
        n: usize,
        h: [usize; n],
	odia: [usize; n],
    }
    let mut dp = vec![0; n];
    for i in 0..n {
        dp[i] = a[i];
        for j in 0..i {
            if h[j] < h[i] {
                dp[i] = max(dp[i], dp[j] + a[i]);
            }    
        }
    }
    let ans = dp.iter().max().unwrap();
    println!("{}", ans);
}

別の解き方

考え方

dp[i][h]を「 $i$ 番目の花までの中で、高さ $h$ の花が一番左側の花だった時の最大の美しさ」として解く。 $i+1$ 番目の高さ $g$ の花を追加した時は、基本的にdp[i+1][h] = dp[i][h]だが、dp[i+1][g]は、 $i+1$ 番目の花だけか、dp[i][h] $(h < g)$ の中で一番美しいものに $i+1$ 番目の花を加えたものになる。dp[i]からdp[i+1]で変化するのは1箇所なので、基本的に同じ行列を使い回せば良い。

ここで後ろの「とある区間で最大値」を求めるのに適しているデータ構造としてSegment Tree, BIT (Bit Intexed Tree）を使うことで、計算量は $O(N\log N)$ にまで落ちる。

BITの実装

まずは、BITを実装する。BITという名前だとビット列を思い浮かべるので（間違ってはいないのだが）、別名である fenwick tree からfenwick_treeという名前で実装した。

fenwick_treeはCrateにも入っているが、区間和に特化しているので今回は使えないのと、
C++のtemplateのようなGenericの練習のために実装してみる。

fenwick_tree.rs

pub struct FenwickTree<T, F> 
{
    table: Vec::<T>,
    op: F,
    default: T
}

impl<T,F> FenwickTree<T,F>
    where 
        T: Clone + Copy,
        F: Fn(T, T) -> T,
{
    pub fn new(n: usize, operation: F, default: T) -> Self {
        FenwickTree { table: vec![default; n + 1], op: operation, default: default}
    }

    pub fn set(&mut self, value: T, mut pos: usize) {
        assert!(pos < self.table.len());
        while pos < self.table.len() {
            self.table[pos] = (self.op)(self.table[pos], value);
            pos += pos & pos.wrapping_neg();
        }
    }

    pub fn get(&self, mut pos: usize) -> T {
        assert!(pos < self.table.len());
        let mut res = self.default;
        while pos != 0 {
            res = (self.op)(res, self.table[pos]);
            pos -= pos & pos.wrapping_neg();
        }
        res
    }
}

Generic

GenericはC++のテンプレートみたいなものだが制約も書ける。ここではTは保持する値の方で、Fは処理する関数である。どんな制約を入れるべきかは、分かりやすいRustのエラーに従えばいい。

impl<T,F> FenwickTree<T,F>
    where 
        T: Clone + Copy,
        F: Fn(T, T) -> T,
{

ちょっと変わっているのが関数の呼び出しで、こんな風にかっこでくくらないといけない。これもコンパイラが教えてくれるが。エラーだけでなく「どうすればいいか」を教えてくれるコンパイラは学習にはいいと思う。

self.table[pos] = (self.op)(self.table[pos], value);

他ファイルの関数の呼び出し

一番シンプルなのはmain.rsと同じ階層にファイルを置いて、

mod fenwick_tree;

と書けばいい。呼び出し方は以下のようにする。maxは std::cmp::maxなので関数。sumにすると区間の和になる。

   let mut dp = fenwick_tree::FenwickTree::new(n, max, 0);

ソースコード

コードとしてはとてもシンプル。
dpの初期値が0なのと、 $d > 0$ から、「dp[i][h] $(h < g)$ の中で一番美しいものに $i+1$ 番目の花を加えたもの」だけを考えていればいい。

main.rs

use proconio::input;
use std::cmp::max;
mod fenwick_tree;

fn main(){
    input! {
        n: usize,
        h: [usize; n],
        a: [usize; n],
    }
    let mut dp = fenwick_tree::FenwickTree::new(n, max, 0);
    for i in 0..n {
        let d: usize = dp.get(h[i]);
        dp.set(d + a[i], h[i]);
    }
    let ans = dp.get(n);
    println!("{}", ans);
}

測定結果

$N$ を $2\times 10^5$ まで変化させた時の処理時間の測定結果。最初の方法では42秒もかかっていた計算が、BITを使うことで6ms以下まで削減できている。