AtCoder Library Practice Contest を ac-library-rs で解いてみた - データ構造編

はじめに

AtCoder 2023年8月の 新ジャッジコンテストから、Rust で ac-library-rs をそのまま use できるようになりました。C++ オリジナルの AtCoder Library の Rust 版です。

• atcoder/ac-library: AtCoder Library
• rust-lang-ja/ac-library-rs: ac-library-rs is a rust port of AtCoder Library (ACL).

これまでは Python スクリプトで切り出して、提出コードに張り付けていました。とても楽になりました。

そこで、 AtCoder Library 公式練習問題を解きつつ、AtCoder Library の使い方を眺めてみようという企画です。

この練習問題は全12問と多いです。本記事ではデータ構造系の 5問を扱います。

本記事で扱うこと

• コードとの対応が分かる程度の問題の解き方図解
  • Rust 以外で解く方にも、競技プログラミングで使うアルゴリズムの参考になるかもしれません
• ac-library-rs の使い方
• Rust + ac-library-rs + proconio, itertools で解いたコード例
  • AtCoder Library や ac-library-rs GitHub 内の解答例を参考にしました:
  • ac-library/document_ja/index.md at master · atcoder/ac-library
    • とても丁寧な公式解説です。C++ 以外の方にもおすすめ
  • ac-library-rs/examples at master · rust-lang-ja/ac-library-rs
  • Add sample code for the practice contest by TonalidadeHidrica · Pull Request #52 · rust-lang-ja/ac-library-rs

本記事で扱わないこと

• AtCoder Library の AtCoder Library Practice Contest では使わない機能については触れません (modint 詳細など)
• AtCoder Library が内部でどのようなデータ構造を持ち、データ処理しているかということは扱いません
• AtCoder Library 以外のライブラリーでも解ける、という話は扱いません (petgraph など) ¹
• 問題の解答方針をどうすれば考え付くかというところは扱いません。練習問題の解説リンクにあるけんちょんさん解説などをどうぞ。

AtCoder Library Practice Contest 問題一覧

• データ構造系 (本記事の対象)
  • B - Fenwick Tree
  • I - Number of Substrings
  • J - Segment Tree
  • K - Range Affine Range Sum
  • L - Lazy Segment Tree
• 数学系
  • C - Floor Sum
  • F - Convolution
• グラフ系
  • A - Disjoint Set Union
  • D - Maxflow
  • E - MinCostFlow
  • G - SCC
  • H - Two SAT

ac-library-rs の使い方

インストール

cargo add ac-library-rs@0.1.1

cargo-compete を使う場合は、新ジャッジ対応で compete.toml を書き換えるついでに、 ac-library-rs に依存関係を付けます。

• qryxip/cargo-compete: A Cargo subcommand for competitive programming
• Release v0.10.6 · qryxip/cargo-compete

[template.new]
dependencies = '''
ac-library-rs = "=0.1.1"
:
'''

コードを書く

    let mut bit = ac_library::FenwickTree::new(n, 0usize);

ac_library::FenwickTree のように crate 名付きで使えます。 use ac_library::FenwickTree; しても良いです。

ac-library-rs の実装の横にテストコードが書いています。こちらを見ると、だいたいの使い方が分かります。

ac-library-rs-0.1.1\src\fenwicktree.rs

    #[test]
    fn fenwick_tree_works() {
        let mut bit = FenwickTree::new(5, 0i64);
        // [1, 2, 3, 4, 5]
        for i in 0..5 {
            bit.add(i, i as i64 + 1);
        }
        assert_eq!(bit.sum(0..5), 15);
        assert_eq!(bit.sum(0..4), 10);
        assert_eq!(bit.sum(1..3), 5);
        // :
    }

使い方の注意などもありますので、先にドキュメント読みをおすすめします。

問題を解く / データ構造

B - Fenwick Tree (BIT)

• B - Fenwick Tree
• atcoder.github.io/ac-library/document_ja/fenwicktree.html

フェニック木 (Binary Indexed Tree) は区間和計算と要素更新に強い構造です。たとえば a[0]+a[1]+...+a[4] が欲しいときに、事前に計算済みの a[0]+a[1]+...+a[3] を使って、2つの値の足し算だけで求められます。

この問題についてはフェニック木をそのまま当てはめられます。

B 解答例

practice2-b.rs

use proconio::input;

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

    let mut bit = ac_library::FenwickTree::new(n, 0usize); // サイズ n のフェニック木作成
    for (i, &x) in a.iter().enumerate() {
        bit.add(i, x); // 初期値 a_i を入れる
    }

    for (kind, u, v) in queries {
        match kind {
            0 => {
                bit.add(u, v); // a_u に v を追加
            }
            1 => println!("{}", bit.sum(u..v)), // 区間和
            _ => unreachable!(),
        }
    }
}

フェニック木を自前で素朴に実装すると「a[1] 開始」「先頭からの和のみ」となります。 [u, v) の区間和は [0, v) の区間和から [0, u) を引く、のように。

AtCoder Library のフェニック木は「a[0] 開始」「任意の区間和」です。扱いやすいです。

セグメント木で解く

フェニック木で解ける問題はセグメント木でも解けます。実装量もあまり変わりません。

セグメント木で解いた差分

practice2-b-segtree.rs

+use ac_library::{Additive, Segtree};
use proconio::input;

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

-   let mut bit = ac_library::FenwickTree::new(n, 0usize);
+   let mut segtree = Segtree::<Additive<usize>>::new(n);
    for (i, &x) in a.iter().enumerate() {
-       bit.add(i, x);
+       segtree.set(i, x);
    }

    for (kind, u, v) in queries {
        match kind {
            0 => {
-               bit.add(u, v);
+               segtree.set(u, segtree.get(u) + v);
            }
-           1 => println!("{}", bit.sum(u..v)),
+           1 => println!("{}", segtree.prod(u..v)),
            _ => unreachable!(),
        }
    }
}

I - Number of Substrings

• I - Number of Substrings
• atcoder.github.io/ac-library/document_ja/string.html

abcbcba の部分文字列は何種類あるか、という問題です。 AtCoder Library の文字列アルゴリズムに、こちらに直接対応するものはありませんが、2つを組み合わせると解けます。

• suffix_array(): 長さ N の文字列に対して、末尾を含む部分文字列 N 個を、昇順に並べて返す
• lcp_array(): 文字列 N 個に対して、隣り合う文字列の先頭が何文字共通しているか (Longest Common Prefix) という、N-1 個の配列を返す

一見、この問題で聞かれている「部分」文字列は「末尾 Suffix」も「先頭 LCP」も関係しません。しかし次のように並び替えると、当てはめ方が見えてきます。

i	suffix	lcp	len 1	2	3	4	5	6	7
0	a	1	a
1	abcbcba	0	a	ab	abc	abcb	abcbc	abcbcb	abcbcba
2	ba	1	b	ba
3	bcba	3	b	bc	bcb	bcba
4	bcbcba	0	b	bc	bcb	bcbc	bcbcb	bcbcba
5	cba	2	c	cb	cba
6	cbcba	-	c	cb	cbc	cbcb	cbcba

長さ 7 の abcbcba の部分文字列は、 7+6+5+4+3+2+1 = 7*8/2 = 28 個あります。そして Suffix array の LCP が 1以上の場合、その数だけ同じ部分文字列になっているということで。合計数から取り除きます。

結果、 28 - (1+0+1+3+0+2) = 21 が得られます。こちらが答えとなります。

I 解答例

practice2-i.rs

use proconio::input;

fn main() {
    input! { s: String }
    let sa = ac_library::suffix_array(s.as_str());
    let la = ac_library::lcp_array(s.as_str(), &sa);
    let answer = (s.len() * (s.len() + 1) / 2) - la.iter().sum::<usize>();
    println!("{answer}");
}

J - Segment Tree

• J - Segment Tree
• atcoder.github.io/ac-library/document_ja/segtree.html

セグメント木は区間の計算に強い構造です。計算方法は「最大値」「加算」などいろいろ設定できます。加算だけのフェニック木より適用できる問題が多く、その分設定が少し難しいです。

セグメント木の設定

設定	フェニック木	セグメント木
要素の単位元	不要 (常に 0)	必要
要素の結合 binary_operation	不要 (常に和)	必要
要素の更新	加算	get, set

設定が難しいですが、「最大値」「加算」というよくあるケースは Segtree::<Max<i32>>::new(n); のように設定が用意されていて、そのまま使えます。 Max<i32> すると次の設定になります。

設定	Segtree::<Max<i32>>	補足
要素の単位元	fn identity() -> i32 { i32::min }	1.max(i32.min) == 1 のように、結合しても値を変えない
要素の結合 binary_operation	fn binary_operation(a: &i32, b: &i32) -> i32 { max(*a, *b) }

この問題では Max<i32> で十分です。

set(), prod()

set() で値の更新を行います。

prod() で Max<i32> による要素の結合、つまり区間の最大値を求めます。

max_right()

「$A_X, ..., A_N$ の中ではじめて $V ≤ A_j$ となる j を求める」という小問もあります。言い換えると「[X, N) の中で、区間 $A_X, ..., A_j$ の最大値がはじめて $V$ 以上となるのはどこか」という問題です。 max_right() で二分探索して解けます。

たとえば i = 1 以上で a[i] >= 3 を探すと次の表のようになります。

i	0	1	2	3	4
A	1	2	3	4	5
区間	[]	[2]	[2, 3]	[2, 3, 4]	[2, 3, 4, 5, 6]
最大値	-	2	3	4	5

答え j = 2 が得られます。

注意点。セグメント木に入れる初期値だけで条件文を満たしてしまうと、二分探索に失敗します。たとえば usize 非負整数に対して a[i] >= 0 を探そうとしてはいけません。 AtCoder Library 公式解説にも書いていました。 ²

f(e()) = true

i32 などの負を扱える型を使う、または単位元について成り立ってしまう場合は max_right() を呼ぶ前に次のコードのように別扱いします。

非負整数で通すコード例 (おすすめしません)

practice2-j-u32.rs

            3 => {
-               input! { x: Usize1, v: i32 }
-               println!("{}", segtree.max_right(x, |a| a < &v) + 1)
+               input! { x: Usize1, v: u32 }
+               let result = if v == 0 {
+                   x + 1 // すべての値は 0 以上なので、調べ始めた index が答え
+               } else {
+                   segtree.max_right(x, |a| a < &v) + 1)
+               };
+               println!("{result}");
            }

J 解答例

practice2-j.rs

use ac_library::{Max, Segtree};
use proconio::{input, marker::Usize1};

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

    let mut segtree = Segtree::<Max<i32>>::new(n);
    for (i, &x) in a.iter().enumerate() {
        segtree.set(i, x);
    }
    for _ in 0..q {
        input! { kind: u8 }
        match kind {
            1 => {
                input! { x: Usize1, v: i32 } // Usize1: 1引いた値を入れる。配列 0 開始にするため
                segtree.set(x, v);
            }
            2 => {
                input! { l: Usize1, r: Usize1 }
                println!("{}", segtree.prod(l..=r));
            }
            3 => {
                input! { x: Usize1, v: i32 }
                println!("{}", segtree.max_right(x, |a| a < &v) + 1)
            }
            _ => unreachable!(),
        }
    }
}

K - Range Affine Range Sum

• K - Range Affine Range Sum
• atcoder.github.io/ac-library/document_ja/lazysegtree.html

「区間の計算」と同時に「区間の更新」も行いたい場合は、遅延セグメント木を使います。普通のセグメント木のように末端の値をそれぞれ更新すると、区間の個数分だけ計算を繰り返し、遅いです。必要になるときまで下側への計算を遅らせる仕組みです。

この問題では「区間和」「区間の更新 $a_{new} = b \times a + c$」を両方行いますので、遅延セグメント木が有効です。

遅延成分の設定

遅延セグメント木の設定は、セグメント木よりも更に多いです。

設定	フェニック木	セグメント木	遅延セグメント木
要素の単位元	不要 (常に 0)	必要	必要
要素の計算 binary_operation	不要 (常に和)	必要	必要
遅延成分の単位元	なし	なし	必要
遅延成分を要素に適用 mapping	なし	なし	必要
遅延成分同士の合成 composition	なし	なし	必要

K 問題については次のように設定します。

セグメント木の区間和なら Segtree::<Additive<usize>> のような仕組みがあります。しかし今回の遅延セグメント木では使えません。区間を更新するために「区間の幅」も欲しいためです。

区間 [2, 4) に 101 を足すときに、区間 [2, 4) の合計値としては 101 × 2 増えます。この計算のため、遅延セグメント木の各要素に「(値, 幅)」 の 2値をセットで覚えるようにします。AtCoder Library の遅延セグメント木では、この幅の持ち方を定石として使うようです。

modint

• atcoder.github.io/ac-library/document_ja/modint.html

modint は数学系ですが、この問題で使っていますので簡単に。

競技プログラミングでは答えとなる整数が 32bit の範囲に収まらない場合のために、「998244353 で割った余りを回答してください。」のような問題が出ることがあります。加減乗算なら計算するたびに余りを求めれば良いという話ですが、計算式が複雑になると忘れてオーバーフローしがちです。また、割り算はひと手間かかります。

ModInt998244353::new(n) のようにして作成した ModInt998244353 型は、この余りの計算を自動で行ってくれます。a0 + a1 のように加減乗除できます。

K 解答例

ac-library-rs の segtree.rs, lazysegtree.rs 使用例をコピーして書き換える感じです。

practice2-k.rs

use ac_library::{LazySegtree, MapMonoid, ModInt998244353 as Mint, Monoid};
use itertools::Itertools;
use proconio::input;

struct Sum;
impl Monoid for Sum {
    type S = (Mint, usize); // (a, size)

    fn identity() -> Self::S {
        (0.into(), 0)
    }

    fn binary_operation(&(a0, n0): &Self::S, &(a1, n1): &Self::S) -> Self::S {
        (a0 + a1, n0 + n1)
    }
}

struct Affine;
impl MapMonoid for Affine {
    type M = Sum;
    type F = (Mint, Mint); // (b, c)

    fn identity_map() -> Self::F {
        (1.into(), 0.into())
    }

    fn mapping(&(b, c): &Self::F, &(a, n): &<Self::M as Monoid>::S) -> <Self::M as Monoid>::S {
        (b * a + c * Mint::new(n), n)
    }

    // b0 * (b1 * a1 + c1) + c0 = (b0 * b1) * a1 + (b0 * c1 + c0)
    fn composition(&(b0, c0): &Self::F, &(b1, c1): &Self::F) -> Self::F {
        (b0 * b1, b0 * c1 + c0)
    }
}

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

    let mut segtree: LazySegtree<Affine> = a.iter().map(|&x| (x, 1usize)).collect_vec().into();
    for _ in 0..q {
        input! { kind: usize }
        match kind {
            0 => {
                input! {
                    l: usize,
                    r: usize,
                    b: Mint,
                    c: Mint,
                }
                segtree.apply_range(l..r, (b, c));
            }
            1 => {
                input! {
                    l: usize,
                    r: usize,
                }
                println!("{}", segtree.prod(l..r).0);
            }
            _ => unreachable!(),
        }
    }
}

L - Lazy Segment Tree

• L - Lazy Segment Tree
• atcoder.github.io/ac-library/document_ja/lazysegtree.html

転倒数を求める問題です。K 問題とおなじ遅延セグメント木ですが、より難しい問題です。

要素と遅延成分の計算式

要素として、その区間内の (0 の個数, 1 の個数, 転倒数) のタプルを持ちます。operation, mapping を DP のように考えて設定します。

この対応関係が分かれば、遅延セグメント木構造につめて終了です。お疲れさまでした。

L 解答例

practice2-l.rs

use ac_library::{LazySegtree, MapMonoid, Monoid};
use itertools::Itertools;
use proconio::{input, marker::Usize1};

struct M;
impl Monoid for M {
    type S = (u64, u64, u64); // #0, #1, #inversion
    fn identity() -> Self::S {
        (0, 0, 0)
    }
    fn binary_operation(&(l0, l1, linv): &Self::S, &(r0, r1, rinv): &Self::S) -> Self::S {
        (l0 + r0, l1 + r1, linv + rinv + l1 * r0)
    }
}

struct F;
impl MapMonoid for F {
    type M = M;
    type F = bool;

    fn identity_map() -> Self::F {
        false
    }
    fn mapping(&f: &Self::F, &(x0, x1, xinv): &<M as Monoid>::S) -> <M as Monoid>::S {
        if f {
            (x1, x0, x0 * x1 - xinv)
        } else {
            (x0, x1, xinv)
        }
    }
    fn composition(&f: &Self::F, &g: &Self::F) -> Self::F {
        f ^ g
    }
}

fn main() {
    input! {
        n: usize,
        q: usize,
        a: [u8; n],
        queries: [(u8, Usize1, Usize1); q],
    }

    let mut segtree: LazySegtree<F> = a
        .iter()
        .map(|s| match s {
            0 => (1, 0, 0),
            1 => (0, 1, 0),
            _ => unreachable!(),
        })
        .collect_vec()
        .into();

    for (t, l, r) in queries {
        match t {
            1 => segtree.apply_range(l..=r, true),
            2 => println!("{}", segtree.prod(l..=r).2),
            _ => unreachable!(),
        }
    }
}

最後に

半月後くらいに「AtCoder Library Practice Contest を ac-library-rs で解いてみた - 数学・グラフ編」が続く予定です。³

AtCoder Library オリジナルの C++ 版と、ac-library-rs など各言語向けに提供している方々に感謝して、本記事を終了します。

1. AtCoder Library は競技プログラミングで良く使うアルゴリズムのうち、C++ の標準ライブラリーには含まれていないものの詰め合わせという感じです。他言語で競技プログラミングする場合は、「AtCoder Library 以外のライブラリーでも解ける」「C++ の標準ライブラリーにあるので AtCoder Library には含まれない」 両方がありますので、AtCoder Library だけ覚えれば大丈夫とはならないと思います。 ↩

2. AtCoder Library 公式解説を読まず、1時間くらいテストを通せずにはまった人がここにいます……。 ↩

3. たぶん AtCoder Library で一番難しい遅延セグメント木を終えたということで、ここまでで満足して力尽きるかもしれません……。 ↩