Rustで遅延評価セグメントツリーを実装してAtCoder Typical90の029を解く。

タイトルの通りです。
遅延評価セグメントツリーを実装します。

遅延評価セグメントツリーがなにかっていう話は他の素敵なブログに譲ります。
この投稿はRustで実装した自分のコードが好きだったのでコレクションとして書いている意図が強いのでご注意ください。

時間をみて詳しい説明や遅延評価セグメントツリー、モノイドなどについての解説を追記するかもしれません。

AtCoderLibraryと同じようにapply(value, left, right)で[left, right)の区間をvalueを使って更新（値を差し替えたり加算したり） する区間更新クエリをさばきprod(left, right)で[left, right)の区間の代表値（最大や和など）を取得する区間取得クエリをさばきます。

pub struct LazySegmentTree<S, F, E, T, G, H, I>
where
    S: Copy + Eq,
    T: Copy + Eq,
    F: Fn(S, S) -> S,
    E: Fn() -> S,
    G: Fn(T, S) -> S,
    H: Fn() -> T,
    I: Fn(T, T) -> T,
{
    value: Vec<S>,
    op: F,
    element: E,
    lazy: Vec<T>,
    mapping: G,
    id: H,
    composite: I,
}

LazySegmentTreeなるStructを実装します。
完全二分木を管理するベクトルvalueと遅延評価を維持するベクトルです。
まずLazyじゃないSegmentTreeにも共通しますが代表値を計算するのに使用する２項演算関数（maxだったりsumだったり）およびその単位元を管理する必要があります。それぞれopとelementに対応します。
続いて評価を遅延するためにためていたlazyを使ってvalueを更新する関数mappingも管理します。
mappingにも単位元が必要で、単位元を出力する関数がidになります。
compositeは合成関数を作る関数です。
mapping(l1, mapping(l2, v)) = mapping(L, v)になるようなLをL = composite(l1, l2)によって計算する関数です。

このStructに区間更新クエリ関数applyと区間取得クエリ関数prodを実装していきましょう。

    pub fn apply(&mut self, v: T, left: usize, right: usize) {
        self._apply(v, 0, left, right, 0, self.value.len() / 2 + 1);
    }

    fn _apply(
        &mut self,
        v: T,
        tree_index: usize,
        search_left: usize,
        search_right: usize,
        left: usize,
        right: usize,
    ) {
        if right <= search_left || search_right <= left {
            return;
        }
        self.lazy[tree_index] = (self.composite)(v, self.lazy[tree_index]);
        if search_left <= left && right <= search_right {
            self.value[tree_index] = (self.mapping)(v, self.value[tree_index]);
        } else {
            self.propagate(tree_index, search_left, search_right, left, right);
            let (left_t_index, right_t_index) = self.get_children(tree_index);
            self.value[tree_index] = (self.op)(self.value[left_t_index], self.value[right_t_index]);
        }
    }

まずはApplyから。
left, rightは現在参照しているtree_indexに相当するノードの子孫の葉ノードが[left, right)であることを示しています。
tree_indexが０（ルート）なら[left, right)はすべての葉ノードを示していますね。

最初のif文は、クエリの支持する更新領域[search_left, search_right)とノードが管理する葉ノードの共通部分が１つもなかったら何もしないというだけの処理ですね。
そうでないときは必ず遅延評価値lazyを新しく入ってきたクエリvによって更新します。

そしてもし現在のノードの子孫の葉ノードがすべて更新領域に含まれる場合には現在のノードの値を更新するだけで処理を終了することで評価を遅延します。
逆に現在のノードの葉ノードのうち更新領域に含まれないものが存在する場合には仕方ないのでlazy値を子孫に伝播します。

    fn propagate(
        &mut self,
        tree_index: usize,
        search_left: usize,
        search_right: usize,
        left: usize,
        right: usize,
    ) {
        let lazy = self.lazy[tree_index];
        self.lazy[tree_index] = (self.id)();
        let mid = (left + right) / 2;
        let (left_t_index, right_t_index) = self.get_children(tree_index);
        self._apply(lazy, left_t_index, search_left, search_right, left, mid);
        self._apply(lazy, right_t_index, search_left, search_right, mid, right);
    }

伝播関数です。self.get_children(tree_index)の機能は察してください。６行もかいていますがやってることはlazy値をid()`に初期化した上で左右のapplyを呼び出すだけですね。

    pub fn prod(&mut self, left: usize, right: usize) -> S {
        assert!(left <= right);
        if right == left {
            (self.element)()
        } else {
            self._prod(0, left, right, 0, self.value.len() / 2 + 1)
        }
    }

    fn _prod(
        &mut self,
        tree_index: usize,
        search_left: usize,
        search_right: usize,
        left: usize,
        right: usize,
    ) -> S {
        if search_left <= left && right <= search_right {
            self.value[tree_index]
        } else if right <= search_left || search_right <= left {
            (self.element)()
        } else {
            if self.lazy[tree_index] != (self.id)() {
                self.propagate(tree_index, left, right, left, right);
            }
            let mid = (left + right) / 2;
            let (left_t_index, right_t_index) = self.get_children(tree_index);
            let l_value = self._prod(left_t_index, search_left, search_right, left, mid);
            let r_value = self._prod(right_t_index, search_left, search_right, mid, right);
            (self.op)(l_value, r_value)
        }
    }

続いて取得クエリです。ちょっと時間ないので詳しくはまた今度書きます。

    pub fn new(size: usize, op: F, element: E, mapping: G, id: H, composite: I) -> Self {
        let tree_size = size.next_power_of_two() * 2 - 1;
        let value = vec![element(); tree_size];
        let lazy = vec![id(); tree_size];
        Self {
            value,
            op,
            element,
            lazy,
            mapping,
            id,
            composite,
        }
    }

コンストラクトです。ちょっと時間がないので（ｒｙ

値の入れ替えによる区間更新と区間最大値取得クエリをさばけるようなLazySegmentTreeでAtCoder典型90の29問目を解いてみよう！

AtCoder 典型９０問　２９問目

問題はこちら

fn main() {
    let (w, n): (usize, usize) = parse_line().unwrap();
    let size = w;
    let op = |x: usize, y: usize| std::cmp::max(x, y);
    let element = || 0usize;
    let mapping = |l: Option<usize>, v: usize| match l {
        Some(val) => val,
        None => v,
    };
    let id = || None::<usize>;
    let composite = |f: Option<usize>, g: Option<usize>| match f {
        Some(_) => f,
        None => g,
    };
    let mut segtree = LazySegmentTree::new(size, op, element, mapping, id, composite);
    let mut ans = Vec::new();
    for _ in 0..n {
        let (l, r): (usize, usize) = parse_line().unwrap();
        let next_height = segtree.prod(l - 1, r) + 1;
        ans.push(next_height);
        segtree.apply(Some(next_height), l - 1, r);
    }
    for val in ans {
        println!("{}", val);
    }
}

このLazySegmentTree,ライブラリ使うのも結構難しいんですよね。
設定しなきゃいけないパラメタが多すぎて、木の大きさはもちろんのことモノイド演算２種とその単位元、および合成関数を作る関数も用意しなきゃいけない。
（設定項目が多いのは僕が悪いわけじゃないはずです！AtCoderLibraryはこれに加えてvalueの型やlazyの型も指定しなきゃいけないですからね！）

さて、問題みていただければわかるようにこの問題は区間[l, r)における最大の高さを取得するクエリと区間[l, r)の値をxに更新するクエリをさばく必要があります。

したがって取得クエリに関連する二項演算は|x, y| std::cmp::maxになります。最大値演算の単位元は値域の最小値です。ここでは自然数を想定してるので0が単位元ですね。

更新クエリは厄介です。
素敵なブログ②を参考にmappingの入力に単位元(id)が来れば値が値を維持し、そうでないときは値を入力で更新するロジックを書きました。
Option型がぴったりで、Noneを単位元として扱うことでスッキリ書くことができましたね。

    let mapping = |l: Option<usize>, v: usize| match l {
        Some(val) => val,
        None => v,
    };

compositeもやってることは一緒です。

    let composite = |f: Option<usize>, g: Option<usize>| match f {
        Some(_) => f,
        None => g,
    };

ということで最終的にこのコードでACできます。

use whiteread::parse_line;

pub struct LazySegmentTree<S, F, E, T, G, H, I>
where
    S: Copy + Eq,
    T: Copy + Eq,
    F: Fn(S, S) -> S,
    E: Fn() -> S,
    G: Fn(T, S) -> S,
    H: Fn() -> T,
    I: Fn(T, T) -> T,
{
    value: Vec<S>,
    op: F,
    element: E,
    lazy: Vec<T>,
    mapping: G,
    id: H,
    composite: I,
}

impl<S, F, E, T, G, H, I> LazySegmentTree<S, F, E, T, G, H, I>
where
    S: Copy + Eq,
    T: Copy + Eq,
    F: Fn(S, S) -> S,
    E: Fn() -> S,
    G: Fn(T, S) -> S,
    H: Fn() -> T,
    I: Fn(T, T) -> T,
{
    pub fn new(size: usize, op: F, element: E, mapping: G, id: H, composite: I) -> Self {
        let tree_size = size.next_power_of_two() * 2 - 1;
        let value = vec![element(); tree_size];
        let lazy = vec![id(); tree_size];
        Self {
            value,
            op,
            element,
            lazy,
            mapping,
            id,
            composite,
        }
    }

    pub fn prod(&mut self, left: usize, right: usize) -> S {
        assert!(left <= right);
        if right == left {
            (self.element)()
        } else {
            self._prod(0, left, right, 0, self.value.len() / 2 + 1)
        }
    }

    fn _prod(
        &mut self,
        tree_index: usize,
        search_left: usize,
        search_right: usize,
        left: usize,
        right: usize,
    ) -> S {
        if search_left <= left && right <= search_right {
            self.value[tree_index]
        } else if right <= search_left || search_right <= left {
            (self.element)()
        } else {
            if self.lazy[tree_index] != (self.id)() {
                self.propagate(tree_index, left, right, left, right);
            }
            let mid = (left + right) / 2;
            let (left_t_index, right_t_index) = self.get_children(tree_index);
            let l_value = self._prod(left_t_index, search_left, search_right, left, mid);
            let r_value = self._prod(right_t_index, search_left, search_right, mid, right);
            (self.op)(l_value, r_value)
        }
    }

    fn get_children(&self, tree_index: usize) -> (usize, usize) {
        (tree_index * 2 + 1, tree_index * 2 + 2)
    }

    pub fn apply(&mut self, v: T, left: usize, right: usize) {
        self._apply(v, 0, left, right, 0, self.value.len() / 2 + 1);
    }

    fn _apply(
        &mut self,
        v: T,
        tree_index: usize,
        search_left: usize,
        search_right: usize,
        left: usize,
        right: usize,
    ) {
        if right <= search_left || search_right <= left {
            return;
        }
        self.lazy[tree_index] = (self.composite)(v, self.lazy[tree_index]);
        if search_left <= left && right <= search_right {
            self.value[tree_index] = (self.mapping)(v, self.value[tree_index]);
        } else {
            self.propagate(tree_index, search_left, search_right, left, right);
            let (left_t_index, right_t_index) = self.get_children(tree_index);
            self.value[tree_index] = (self.op)(self.value[left_t_index], self.value[right_t_index]);
        }
    }

    fn propagate(
        &mut self,
        tree_index: usize,
        search_left: usize,
        search_right: usize,
        left: usize,
        right: usize,
    ) {
        let lazy = self.lazy[tree_index];
        self.lazy[tree_index] = (self.id)();
        let mid = (left + right) / 2;
        let (left_t_index, right_t_index) = self.get_children(tree_index);
        self._apply(lazy, left_t_index, search_left, search_right, left, mid);
        self._apply(lazy, right_t_index, search_left, search_right, mid, right);
    }
}

fn main() {
    let (w, n): (usize, usize) = parse_line().unwrap();
    let size = w;
    let op = |x: usize, y: usize| std::cmp::max(x, y);
    let element = || 0usize;
    let mapping = |l: Option<usize>, v: usize| match l {
        Some(val) => val,
        None => v,
    };
    let id = || None::<usize>;
    let composite = |f: Option<usize>, g: Option<usize>| match f {
        Some(_) => f,
        None => g,
    };
    let mut segtree = LazySegmentTree::new(size, op, element, mapping, id, composite);
    let mut ans = Vec::new();
    for _ in 0..n {
        let (l, r): (usize, usize) = parse_line().unwrap();
        let next_height = segtree.prod(l - 1, r) + 1;
        ans.push(next_height);
        segtree.apply(Some(next_height), l - 1, r);
    }
    for val in ans {
        println!("{}", val);
    }
}

今回は以上です。雑い自己満記事ですが付き合っていただきありがとうございました。