はじめに

競技プログラミングで、しゃくとり法 (Two-pointer) を使いたくなることがあります。

その場で短く実装しようとすると、いろいろな操作を 1つのループ中に書きがちです。うっかりバグらせてしまうことも多いです。

そこで、先人の実装を参考に、Rust で動くしゃくとり法ライブラリを作ってみました。

このライブラリが実用的かどうかは分かりません。「しゃくとり法を部品に分けて考える」ことの練習になるはずです。そんな紹介記事です。

2024年エイプリルフールのネタ記事でもあります。

本記事で行うこと

しゃくとり法の簡単な紹介と、複数の関数に分ける考え方を示します
Rust のイテレーターで扱うライブラリを 3通り実装します
- 「A. 有効な区間のみ返す」Iterator トレイトを満たす実装
- 「B. 有効な区間と内部状態のコピーを返す」 Iterator トレイトを満たす実装
- 「C. 有効な区間と内部状態の参照を返す」 Rust 1.65 GATs の LendingIterator 例を満たす実装
しゃくとり法の例題 6問を、A, B, C と「D. しゃくとり法ライブラリなし」の 4通りで解きます
1. 単調増加な数列から 2 つの値を選び、差が k になるものがあるかを求める
2. 数列の区間のうち、合計値が k になるものがあるかを求める
3. 数列の区間のうち、合計値が 10 以下になる個数を求める
4. 数列の区間のうち、単調増加となる区間の個数を求める (ABC038-C)
5. 数列の区間のうち、最大値と最小値の差が k 以下となる最大の長さを求める
6. 数列の区間のうち、現れる数字が k 種類以下となる最大の長さを求める (典型90-034)

本記事で行わないこと

しゃくとり以外の Two-pointer は、本記事のライブラリでは扱いません。たとえば以下のものは対象外です
- 両端から区間を縮めるもの
- 右端にたどり着いたあと、ループして左端に戻るもの
- 左端が右端を追い越すもの
他の手法と組み合わせてしゃくとり法を複数回実行する問題は扱いません

しゃくとり法の考え方

配列 [1, 3, 3, 4, 5] の中に、数の合計が 10 となる区間があるか、という問題を考えます。

手で解く

選択している区間の左端と右端を、しゃくとり虫のように動かして解きます。

区間の左と右を差す変数 l, r を設定する。これらを動かし、合計 sum を更新する
sum = 10: 見つかった → 終了
sum < 10: r を増やす
sum > 10: l を増やす

[0]	[1]	[2]	[3]	l	r	sum	check	next
1				0	1	1	< 10	r += 1
1	3			0	2	4	< 10	r += 1
1	3	3		0	3	7	< 10	r += 1
1	3	3	4	0	4	11	> 10	l += 1
	3	3	4	1	4	10	= 10

途中で [3, 3, 4] の組が見つかりました。ここで終了です。

区間の左端と右端という 2つの場所を指していることから、しゃくとり法を Two-pointer と呼ぶこともあります。

しゃくとり法をより詳しく知りたい方は、けんちょんさんの解説記事をどうぞ:

しゃくとり法 (尺取り法) の解説と、それを用いる問題のまとめ #競技プログラミング - Qiita

単体テスト

次のテストが通るような関数 find_sum10() をこれから考えます。

#[test]
fn test_find_sum10() {
    assert_eq!(find_sum10(&[1, 3, 3, 4, 5]), true); // 3 + 3 + 4 = 10
    assert_eq!(find_sum10(&[1, 2, 3, 4, 5]), true); // 1 + 2 + 3 + 4 = 10
    assert_eq!(find_sum10(&[1, 3, 0, 4, 5]), false);
    assert_eq!(find_sum10(&[1, 100, 1, 9, 1]), true);  // 1 + 9 = 10
    assert_eq!(find_sum10(&[]), false);
}

as-is: 1つの関数で実装する例

しゃくとり法の実装の一つとして、次のように短く書けます。for ループ内は 10行程度です。

fn find_sum10(a: &[usize]) -> bool {
    let n = a.len();

    let mut l = 0;
    let mut s = 0; // 状態: 区間の合計
    for r in 0..n { // 右端を動かす
        s += a[r]; // 状態の更新
        while l <= r && s > 10 { // 左端を動かすかの判定
            s -= a[l]; // 状態の更新
            l += 1; // 左端を動かす
        }
        if s == 10 { // 状態の判定
            return true;
        }
    }

    false
}

さて、この関数は難しくないですが、更新しやすい形でもありません。

左端を動かす条件 s > 10 と、探したい条件 s == 10 が離れた場所に書かれています
右端を動かしたときの状態更新 s += a[r]; と、左端を動かしたの状態更新 s -= a[l]; の処理が違うインデントにあります。対応関係が分かりにくいです

また、しゃくとり法の実装方法は多数あります。それぞれの変数が何を指すかが参考実装によって異なりがちで、問題にあわせて実装を更新する際に混乱しやすいです。たとえば次のように。

区間 l, r の扱い
- 半開区間 [l, r) で扱う派 (本記事のイテレーター系の実装)
- 閉区間 [l, r] で扱う派 (本記事の他の実装)
l か r を追い越さないようにする処理の書き方
- while ループの条件に l < r などを書き、追い越しを防ぐ派 (本記事のイテレーター系の実装)
- 区間の判定で自然に l < r などが満たされる問題については書かない派 (本記事の他の実装)
ループの回し方
- 2重ループ: l は [0, n) の範囲内を順に動かし、そのループ中で r を [l, len] の範囲内で条件を満たすよう while 文で更新する派
- 2重ループ: r は [0, n) の範囲内を順に動かし、そのループ中で l を [0, r] の範囲内で条件を満たすよう while 文で更新する派 (本記事の実装)
- 1重ループ: while ループの中で、区間が条件を満たせば r += 1, 満たさないときに l += 1 する操作を繰り返す派
結果 s == 10 の判定呼び出しタイミング
- 外側のループに対して 1回だけ判定する派 (本記事の実装)
- l または r が 1つ動くたびに判定する派

こうパターンがいろいろとあると、参考実装があってもどの流れか、どう書き換えれば良いかの把握に時間がかかってしまいます。それぞれ一長一短があり、決定打となるものはありません。 ¹

to-be: イテレーターで解く例

「B. 有効な区間と内部状態のコピーを返す」イテレーターを使って実装してみます。

イテレーターが (l, r, sum) に有効な範囲と状態 (0, 1, 1), (0, 2, 4), ... を順に返してくれるものとします。sum は 10以下になる範囲で、l を最大まで伸ばしたものです。

[0]	[1]	[2]	[3]	l	r	sum	check
1				0	1	1	< 10
1	3			0	2	4	< 10
1	3	3		0	3	7	< 10
	3	3	4	1	4	10	= 10

fn find_sum10(a: &[usize]) -> bool {
    TwoPointerStateIter::<Sum10Ops>::new(a, 0).any(|(l, r, s)| s == 10)
}

そうすると、イテレーター TwoPointerStateIter を利用する側としては、合計値 sum == 10 かどうかを判定すれば終了です。問題が解けました。

イテレーターを動かすための演算方法 Sum10Ops は別に設定します。 AtCoder Library でセグメント木や遅延セグメント木の処理方法を外部から指定するのと同じ感じです。

struct Sum10Ops;
impl TwoPointerIterOps for Sum10Ops {
    type I = usize; // 配列の値の型
    type S = usize; // 状態の型

    fn test(&state: &Self::S, a: &[Self::I]) -> bool {
        state <= 10 // 選択範囲が条件を満たしているか
    }
    fn on_push_back(state: &mut Self::S, _: usize, x: &Self::I) {
        *state += x; // 右端を動かしたときの状態更新
    }
    fn on_pop_front(state: &mut Self::S, _: usize, x: &Self::I) {
        *state -= x; // 左端を動かしたときの状態更新
    }
}

とても更新箇所がスッキリします。このようにしゃくとり法を扱いたい、という記事です。

フローチャートを描く

フローチャートでしゃくとり法を描くと、次のイメージです。

イテレーターを使う場合は、外側のループ while r < n に対応するところが、1つずつ値を取り出して外側で処理できるようになるイメージです。

Rust のイテレーターで扱うライブラリ実装

3通りのイテレーターと、それらで使う共通の演算方法を実装します。

方法	利点 Pros.	欠点 Cons.
A. 有効な区間のみ返す	Iterator のメソッドを自由に使える	内部状態が分からない代わりに呼び出し側の前準備が少し増える
B. 有効な区間と内部状態のコピーを返す	Iterator のメソッドを自由に使える	オブジェクトのコピーが重く使えない場合あり
C. 有効な区間と内部状態の参照を返す	オブジェクトも参照できる	Iterator のメソッドが使えない

イテレーターの実装は、あずんひさんのスライドを参考にしました。いつもお世話になっています。

Rustのイテレーター完全制覇 / domination-of-the-rust-iterators - Speaker Deck

3種類のイテレーター利用例

先ほどと同じ、配列 [1, 3, 3, 4, 5] の中に、数の合計が 10 となる区間があるか、という問題を考えます。

3種類のイテレーターはすべて、 r を 1つずつ順に増やし、l は区間の合計が 10 を超えない左端を返します。配列の要素数だけ値を返します。

[0]	[1]	[2]	[3]	[4]	l	r	sum	A. next()	B. next()	C. next()
1					0	1	1	(0, 1)	(0, 1, 1)	(0, 1, &1)
1	3				0	2	4	(0, 2)	(0, 2, 4)	(0, 2, &4)
1	3	3			0	3	7	(0, 3)	(0, 3, 7)	(0, 3, &7)
	3	3	4		1	4	10	(1, 4)	(1, 4, 10)	(1, 4, &10)
			4	5	3	5	9	(3, 5)	(3, 5, 9)	(3, 5, &9)

ライブラリ利用側は、区間に対して数の合計が 10 かどうかだけを調べれば良いです。

利用コードは次のようになります:

fn find_sum10_a(a: &[usize], cum: &[usize]) -> bool {
    TwoPointerIter::<Sum10Ops>::from(a).any(|(l, r)| cum[r] - cum[l] == 10)
}

fn find_sum10_b(a: &[usize]) -> bool {
    TwoPointerStateIter::<Sum10Ops>::from(a).any(|(_, _, state)| state == 10)
}

fn find_sum10_c(a: &[usize]) -> bool {
    let mut iter = TwoPointerLendingIter::<Sum10Ops>::from(a);
    while let Some((_, _, &state)) = iter.next() {
        if state == 10 {
            return true;
        }
    }
    false
}

A: 内部状態が利用者からは分かりません。利用者も区間の合計が素早く求められるように、事前に累積和を用意する、というひと手間がかかります。
B: 内部状態をコピーで取得しています。
C: 内部状態を参照しています。不便ですが、コピーが重いときに利用できます。

本記事で扱う例題では、すべて「A. 有効な区間のみ返す」で扱えます。この問題のように「B. 有効な区間と内部状態のコピーを返す」を使うと楽になることもあります。

演算方法

演算方法 TwoPointerIterOps です。 3通りのイテレーターで共通して使います。

trait TwoPointerIterOps {
    /// Item of array
    type I;

    /// State
    type S: Clone + Default;

    /// 選択区間が条件を満たしているかを返す
    ///
    /// # Arguments
    ///
    /// * `state` - 状態
    /// * `a` - 配列の選択区間 (空にはならない)
    fn test(state: &Self::S, a: &[Self::I]) -> bool;

    /// 選択区間の右側を 1つ削除するときのコールバック。状態を更新する際に設定する
    ///
    /// # Arguments
    ///
    /// * `state` - 状態
    /// * `i` - 追加する位置
    /// * `x` - 追加する値 (== a[i])
    #[allow(unused_variables)]
    fn on_push_back(state: &mut Self::S, i: usize, x: &Self::I) {}

    /// 選択区間の左側を 1つ削除するときのコールバック。状態を更新する際に設定する
    ///
    /// # Arguments
    ///
    /// * `state` - 状態
    /// * `i` - 削除する位置
    /// * `x` - 削除する値 (== a[i])
    #[allow(unused_variables)]
    fn on_pop_front(state: &mut Self::S, i: usize, x: &Self::I) {}
}

各関数の引数は、よく使いそうなものに絞っています。

演算方法でどこまで状態を公開するかについて

いろいろな状態を使えるようにしようと思うと、次のような形も考えられます。 a は配列全体を指します。

trait TwoPointerPowerfulIterOps {
    type I;
    type S: Clone + Default;

    fn test(state: &Self::S, a: &[Self::I], l: usize, r: usize) -> bool;
    fn on_push_back(state: &mut Self::S, a: &[Self::I], l: usize, r: usize);
    fn on_pop_front(state: &mut Self::S, a: &[Self::I], l: usize, r: usize);
}

しかし、参照できる状態が多いほど、使い方に困ってしまいます。右に伸ばすときにも l が欲しいことはまれです。まれなら、見えない方が分かりやすいです。

それに、l <= r か、a[l] や a[r] は範囲外アクセスにならないかというところについて、ライブラリの利用者が意識したくないところです。使い勝手が良くないです。

そのため、本ライブラリでは良く使いそうなものだけを引数で渡すようにしています。

    fn test(state: &Self::S, a: &[Self::I]) -> bool;
    fn on_push_back(state: &mut Self::S, i: usize, x: &Self::I);
    fn on_pop_front(state: &mut Self::S, i: usize, x: &Self::I);

update 系のインデックス i も使う機会は少ないです。 for (i, &x) in a.iter().enumerate() との対応から、まああっても良いか、くらいの感じです。例題 5. で使用します。

A. 有効な区間のみ返すイテレーター

ここからはしばらく、ライブラリ利用者は意識しなくて良い話が続きます。使い方を知りたい方は例題まで進んでください。

struct TwoPointerIter<'a, Ops>
where
    Ops: TwoPointerIterOps,
{
    a: &'a [Ops::I],
    l: usize,
    r: usize,
    state: Ops::S,
}

impl<'a, Ops: TwoPointerIterOps> TwoPointerIter<'a, Ops> {
    fn new(a: &'a [Ops::I], state: Ops::S) -> Self {
        let (l, r) = (0, 0);
        TwoPointerIter::<'a, Ops> { a, l, r, state }
    }
}

impl<'a, Ops: TwoPointerIterOps> From<&'a [Ops::I]> for TwoPointerIter<'a, Ops> {
    fn from(a: &'a [Ops::I]) -> Self {
        Self::new(a, Ops::S::default())
    }
}

impl<'a, Ops: TwoPointerIterOps> Iterator for TwoPointerIter<'a, Ops> {
    type Item = (usize, usize);

    fn next(&mut self) -> Option<Self::Item> {
        let a = self.a;
        if self.r < a.len() {
            Ops::on_push_back(&mut self.state, self.r, &a[self.r]);
            self.r += 1;

            while self.l < self.r && !Ops::test(&self.state, &a[self.l..self.r]) {
                Ops::on_pop_front(&mut self.state, self.l, &a[self.l]);
                self.l += 1;
            }

            Some((self.l, self.r))
        } else {
            None
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.a.len(), Some(self.a.len()))
    }
}

impl<'a, Ops: TwoPointerIterOps> ExactSizeIterator for TwoPointerIter<'a, Ops> {
    fn len(&self) -> usize {
        self.a.len()
    }
}

impl<'a, Ops: TwoPointerIterOps> std::iter::FusedIterator for TwoPointerIter<'a, Ops> {}

以下を実装しました:


TwoPointerIter 構造体	配列と状態に加え、区間を表す `[l, r)` を持つ
TwoPointerIter.new()	配列と初期状態を受け取る。初期区間は `[0, 0)`
TwoPointerIter.from()	配列を受け取る。初期状態は型のデフォルト値。初期区間は `[0, 0)`
TwoPointerIter.next()	次の有効な区間 (空を含む) `[l, r)` を返す
TwoPointerIter.size_hint()
TwoPointerIter.len()
FusedIterator マーカー

A. の内部状態

[0]	[1]	[2]	[3]	l	r	state	return
				0	0	0
1				0	1	0 + 1 = 1	(0, 1)
1	3			0	2	1 + 3 = 4	(0, 2)
1	3	3		0	3	4 + 3 = 7	(0, 3)
1	3	3	4	0	4	7 + 4 = 11
	3	3	4	1	4	11 - 1 = 10	(1, 4)

r を右に動かすと、場合によっては一時的に、区間の状態が有効でなくなります。ライブラリ内で直後に l を右に動かすことで、有効な状態に戻ります。ライブラリの利用者は有効な状態以外を知らなくても良いです。

B. 有効な区間と内部状態のコピーを返すイテレーター

struct TwoPointerStateIter<'a, Ops>
where
    Ops: TwoPointerIterOps,
{
    a: &'a [Ops::I],
    l: usize,
    r: usize,
    state: Ops::S,
}

impl<'a, Ops: TwoPointerIterOps> TwoPointerStateIter<'a, Ops> {
    fn new(a: &'a [Ops::I], state: Ops::S) -> Self {
        let (l, r) = (0, 0);
        TwoPointerStateIter::<'a, Ops> { a, l, r, state }
    }
}

impl<'a, Ops: TwoPointerIterOps> From<&'a [Ops::I]> for TwoPointerStateIter<'a, Ops> {
    fn from(a: &'a [Ops::I]) -> Self {
        Self::new(a, Ops::S::default())
    }
}

impl<'a, Ops: TwoPointerIterOps> Iterator for TwoPointerStateIter<'a, Ops> {
    type Item = (usize, usize, Ops::S);

    fn next(&mut self) -> Option<Self::Item> {
        let a = self.a;
        if self.r < a.len() {
            Ops::on_push_back(&mut self.state, self.r, &a[self.r]);
            self.r += 1;

            while self.l < self.r && !Ops::test(&self.state, &a[self.l..self.r]) {
                Ops::on_pop_front(&mut self.state, self.l, &a[self.l]);
                self.l += 1;
            }

            Some((self.l, self.r, self.state.clone()))
        } else {
            None
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.a.len(), Some(self.a.len()))
    }
}

impl<'a, Ops: TwoPointerIterOps> ExactSizeIterator for TwoPointerStateIter<'a, Ops> {
    fn len(&self) -> usize {
        self.a.len()
    }
}

impl<'a, Ops: TwoPointerIterOps> std::iter::FusedIterator for TwoPointerStateIter<'a, Ops> {}

A. 実装からの変更点は、主に下記 2行です:

rust

-    type Item = (usize, usize);
+    type Item = (usize, usize, Ops::S);

rust

-            Some((self.l, self.r)
+            Some((self.l, self.r, self.state.clone()))

内部状態 self.state の型は、clone() でコピーできるよう Clone トレイトを満たすようにしています。 usize のような単純な数値も、 HashMap のような複雑な状態も、clone() できるならライブラリ側としては同じように扱えます。

もちろん、HashMap のような複雑な状態を持つと、clone() も重いです。そんな処理が iter() を呼ぶたびに毎回呼ばれてしまいます。これだけでタイムアウトしかねません。

内部状態が複雑な場合は、この方法ではなく A. や C. で行いたいです。

C. 有効な区間と内部状態の参照を返すイテレーター

struct TwoPointerLendingIter<'a, Ops>
where
    Ops: TwoPointerIterOps,
{
    a: &'a [Ops::I],
    l: usize,
    r: usize,
    state: Ops::S,
}

impl<'a, Ops: TwoPointerIterOps> TwoPointerLendingIter<'a, Ops> {
    fn new(a: &'a [Ops::I], state: Ops::S) -> Self {
        let (l, r) = (0, 0);
        TwoPointerLendingIter::<'a, Ops> { a, l, r, state }
    }
}

impl<'a, Ops: TwoPointerIterOps> From<&'a [Ops::I]> for TwoPointerLendingIter<'a, Ops> {
    fn from(a: &'a [Ops::I]) -> Self {
        Self::new(a, Ops::S::default())
    }
}

trait LendingIterator {
    type Item<'s>
    where
        Self: 's;
    fn next<'s>(&'s mut self) -> Option<Self::Item<'s>>;
}

impl<'a, Ops: TwoPointerIterOps> LendingIterator for TwoPointerLendingIter<'a, Ops> {
    type Item<'s> = (usize, usize, &'s Ops::S) where Self: 's;

    fn next<'s>(&'s mut self) -> Option<Self::Item<'s>> {
        let a = self.a;
        if self.r < a.len() {
            Ops::on_push_back(&mut self.state, self.r, &a[self.r]);
            self.r += 1;

            while self.l < self.r && !Ops::test(&self.state, &a[self.l..self.r]) {
                Ops::on_pop_front(&mut self.state, self.l, &a[self.l]);
                self.l += 1;
            }

            Some((self.l, self.r, &self.state))
        } else {
            None
        }
    }
}

B. 実装と C. 実装の差分

rust

-struct TwoPointerStateIter<'a, Ops>
+struct TwoPointerLendingIter<'a, Ops>
 where
     Ops: TwoPointerIterOps,
 {
     a: &'a [Ops::I],
     l: usize,
     r: usize,
     state: Ops::S,
 }

-impl<'a, Ops: TwoPointerIterOps> TwoPointerStateIter<'a, Ops> {
+impl<'a, Ops: TwoPointerIterOps> TwoPointerLendingIter<'a, Ops> {
     fn new(a: &'a [Ops::I], state: Ops::S) -> Self {
         let (l, r) = (0, 0);
-        TwoPointerStateIter::<'a, Ops> { a, l, r, state }
+        TwoPointerLendingIter::<'a, Ops> { a, l, r, state }
     }
 }
 
-impl<'a, Ops: TwoPointerIterOps> From<&'a [Ops::I]> for TwoPointerStateIter<'a, Ops> {
+impl<'a, Ops: TwoPointerIterOps> From<&'a [Ops::I]> for TwoPointerLendingIter<'a, Ops> {
     fn from(a: &'a [Ops::I]) -> Self {
         Self::new(a, Ops::S::default())
     }
 }
+
+trait LendingIterator {
+    type Item<'s>
+    where
+        Self: 's;
+    fn next<'s>(&'s mut self) -> Option<Self::Item<'s>>;
+}

-impl<'a, Ops: TwoPointerIterOps> Iterator for TwoPointerStateIter<'a, Ops> {
-    type Item = (usize, usize, Ops::S);
+impl<'a, Ops: TwoPointerIterOps> LendingIterator for TwoPointerLendingIter<'a, Ops> {
+    type Item<'s> = (usize, usize, &'s Ops::S) where Self: 's;
 
-    fn next(&mut self) -> Option<Self::Item> {
+    fn next<'s>(&'s mut self) -> Option<Self::Item<'s>> {
         let a = self.a;
         if self.r < a.len() {
             Ops::on_push_back(&mut self.state, self.r, &a[self.r]);
             self.r += 1;
 
             while self.l < self.r && !Ops::test(&self.state, &a[self.l..self.r]) {
                 Ops::on_pop_front(&mut self.state, self.l, &a[self.l]);
                 self.l += 1;
             }
 
-            Some((self.l, self.r, self.state.clone()))
+            Some((self.l, self.r, &self.state))
         } else {
             None
         }
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
         (self.a.len(), Some(self.a.len()))
     }
 }
-
-impl<'a, Ops: TwoPointerIterOps> ExactSizeIterator for TwoPointerStateIter<'a, Ops> {
-    fn len(&self) -> usize {
-        self.a.len()
-    }
-}
-
-impl<'a, Ops: TwoPointerIterOps> std::iter::FusedIterator for TwoPointerStateIter<'a, Ops> {}

以下を実装しました:


TwoPointerLendingIter 構造体	配列と状態に加え、区間を表す `[l, r)` を持つ
TwoPointerLendingIter.new()	配列と初期状態を受け取る。初期区間は `[0, 0)`
TwoPointerLendingIter.from()	配列を受け取る。初期状態は型のデフォルト値。初期区間は `[0, 0)`
TwoPointerLendingIter.next()	次の有効な区間 (空を含む) `[l, r)` と状態の参照を返す

イテレーターの変化する状態を参照で返すことは難しい

B. の実装はコピーが重いことが問題でした。

コピーの代わりに参照を返すようにすればうまくいくようにも思います。区間が有効かの確認のために、値を書き換えたくなる機会は少ないです。状態を覗き見ることができれば十分なはずです。

rust

             if l < r - 1 {
-                return Some((l, r - 1, self.state_out.clone()));
+                return Some((l, r - 1, &self.state_out));
             }

しかし残念ながら、この方針はうまくいきません。

コピーは iter.next() の返値の寿命が、イテレーターと切り離されます。いつまででも使えます。そのため、複数の状態を windows() などで組み合わせて使うことができます。

let mut iter = TwoPointerStateIter::<Sum10Ops>::new(&[1, 3, 3, 4, 5], 0);
let x0 = iter.next();
let x1 = iter.next();
assert_eq!(x0, Some((0, 3, 7)));
assert_eq!(x1, Some((1, 4, 10)));

参照にすると、次に next() を呼んだときに内部状態が書き変わってしまいます。

let mut iter = TwoPointerBadLendingIter::<Sum10Ops>::new(&[1, 3, 3, 4, 5], 0);
let x0 = iter.next(); // state の参照先が 7 
let x1 = iter.next(); // state の参照先が 10 。x0 の参照先が書き変わってしまっている？

このような参照は、Rust の所有権システムでは認められません。コンパイルエラーになります。Iterator トレイトを満たせません。

しゃくとり法では次のように any() の中でだけ状態を参照できれば十分、このように使い手が注意すれば参照先が書き変わらない範囲で安全に使える、と言いたいところです。でも Iterator では行えません。

fn find_sum10_c_compile_error(a: &[usize]) -> bool {
    TwoPointerLendingIter::<Sum10Ops>::from(a).any(|(_, _, &state)| state == 10)
}

状態を貸し出せる別の仕組みが欲しくなります。

GATs で寿命を制限すれば参照を返せる

このような場合で、Rust 1.65 で追加された Generic associated types (GATs) が使えます。 next() 関数の返値に特定のライフタイムを割り当てられます。

Generic associated types to be stable in Rust 1.65 | Rust Blog

GATs が有効な例として、LendingIterator 状態を貸し出すイテレーターが書かれています。まさに、本記事で欲しいと思っていたものです。

let mut iter = TwoPointerLendingIter::<Sum10Ops>::new(&[1, 3, 3, 4, 5], 0);
assert_eq!(iter.next(), Some((0, 3, &7))); // state 参照先は 7 
assert_eq!(iter.next(), Some((1, 4, &10))); // state 参照先が 10

はい、状態を参照できる形で next() できるようになりました。祝。

もちろん、誰かが参照中に next() で状態を更新しようとするとコンパイルエラーになります。Rust の所有権の原則通りです。

let mut iter = TwoPointerLendingIter::<Sum10Ops>::new(&[1, 3, 3, 4, 5], 0);
let x0 = iter.next(); // state 参照先は 7 
let x1 = iter.next(); // x0 が参照中の値を書き換えようとしているためコンパイルエラー

さて、残念な点があります。これはイテレーターと言う名前ですが Iterator トレイトとは別物です。 any() などの Iterator の機能は、現時点ではなにも使えません ² 。

while ループで値を順に取り出して使います。A., B. に比べて呼び出しが不便、というところが欠点です。

fn find_sum10_c(a: &[usize]) -> bool {
    let mut iter = TwoPointerLendingIter::<Sum10Ops>::from(a);
    while let Some((_, _, &state)) = iter.next() {
        if state == 10 {
            return true;
        }
    }
    false
}

しゃくとり法の例題 6問を解いてみる

しゃくとり法の例題 6問を、A, B, C と「D. しゃくとり法ライブラリなし」の 4通りで解きます。

最初の 1問は「2つの値を選ぶ」、次の 5問は「区間を探す」ものです。

1. 単調増加な数列から 2 つの値を選び、差が `k` になるものがあるかを求める

単体テスト:

#[test]
fn test_find_sub() {
    assert_eq!(find_sub(&[0, 1, 3, 6, 10, 15], 10), true); // 10 - 0 = 10
    assert_eq!(find_sub(&[0, 1, 3, 6, 10, 15], 9), true); // 10 - 1 = 9
    assert_eq!(find_sub(&[0, 1, 3, 6, 10, 15], 8), false);
    assert_eq!(find_sub(&[0, 1, 101, 102, 111, 112], 11), true); // 112 - 101 = 11
    assert_eq!(find_sub(&[], 1), false);
    assert_eq!(find_sub(&vec![10; 10000], 10), false);
}

制約: k > 0

演算方法:

struct SubOps;
impl TwoPointerIterOps for SubOps {
    type I = usize;
    type S = usize;

    fn test(&state: &Self::S, a: &[Self::I]) -> bool {
        a[a.len() - 1] - a[0] <= state
    }
}

状態: state == k 区間の判定に使う固定値
- on_push_back(), on_pop_front() は状態固定のため設定不要

A. 有効な区間のみ返すイテレーター:

fn find_sub_a(a: &[usize], k: usize) -> bool {
    TwoPointerIter::<SubOps>::new(a, k).any(|(l, r)| a[r - 1] - a[l] == k)
}

B. 有効な区間と内部状態のコピーを返すイテレーター:

fn find_sub_b(a: &[usize], k: usize) -> bool {
    TwoPointerStateIter::<SubOps>::new(a, k).any(|(l, r, _)| a[r - 1] - a[l] == k)
}

C. 有効な区間と内部状態の参照を返すイテレーター

fn find_sub_c(a: &[usize], k: usize) -> bool {
    let mut iter = TwoPointerLendingIter::<SubOps>::new(a, k);
    while let Some((l, r, _)) = iter.next() {
        if a[r - 1] - a[l] == k {
            return true;
        }
    }
    false
}

D. しゃくとり法ライブラリなし

fn find_sub_d(a: &[usize], k: usize) -> bool {
    let n = a.len();

    let mut l = 0;
    for r in 0..n {
        while a[r] - a[l] > k {
            l += 1;
        }
        if a[r] - a[l] == k {
            return true;
        }
    }

    false
}

状態は k 固定です。イテレーター内部で、区間が有効かどうかの判定にだけ使います。

有効な区間を取り出したあと、その区間の中に a[r - 1] - a[l] == k なものが含まれているかを、イテレーターの any() メソッドで調べます。

状態更新の例:

assert_eq!(find_sub(&[0, 1, 3, 6, 10, 15], 10), true); // 10 - 0 = 10

[0]	[1]	[2]	[3]	[4]	state	l	r	check
0					9	0	1	0 - 0 < 9³
0	1				9	0	2	1 - 0 < 9
0	1	3			9	0	3	3 - 0 < 9
0	1	3	6		9	0	4	6 - 0 < 9
0	1	3	6	10	9	0	5
	1	3	6	10	9	1	5	10 - 1 = 9

状態を取り出して使うことはありません。シンプルな「A. 有効な区間のみ返すイテレーター」で十分です。

2. 数列の区間のうち、合計値が `k` になるものがあるかを求める

単体テスト:

#[test]
fn test_find_sum() {
    assert_eq!(find_sum(&[1, 2, 3, 4, 5], 10), true); // [1, 2, 3, 4].sum() = 10
    assert_eq!(find_sum(&[1, 2, 3, 4, 5], 9), true); // [2, 3, 4].sum() = 9
    assert_eq!(find_sum(&[1, 2, 3, 4, 5], 8), false);
    assert_eq!(find_sum(&[1, 100, 1, 9, 1], 11), true); // [1, 9, 1].sum() = 11
    assert_eq!(find_sum(&[], 1), false);
    assert_eq!(find_sum(&vec![0; 10000], 1), false);
}

演算方法:

struct SumOps;
impl TwoPointerIterOps for SumOps {
    type I = usize;
    type S = (usize, usize);

    fn test(&(sum, k): &Self::S, _: &[Self::I]) -> bool {
        sum <= k
    }
    fn on_push_back((sum, _): &mut Self::S, _: usize, &x: &Self::I) {
        *sum += x;
    }
    fn on_pop_front((sum, _): &mut Self::S, _: usize, &x: &Self::I) {
        *sum -= x;
    }
}

状態: (sum, k)
- sum: 区間の合計
- k: 区間の判定に使う固定値

A. 有効な区間のみ返すイテレーター:

fn find_sum_a(a: &[usize], k: usize) -> bool {
    let mut cum = vec![0; a.len() + 1]; // 累積和
    for (i, &x) in a.iter().enumerate() {
        cum[i + 1] = cum[i] + x;
    }
    TwoPointerIter::<SumOps>::new(a, (0, k)).any(|(l, r)| cum[r] - cum[l] == k)
    // TwoPointerIter::<SubOps>::new(&cum, k).any(|(l, r)| cum[r - 1] - cum[l] == k)
}

B. 有効な区間と内部状態のコピーを返すイテレーター:

fn find_sum_b(a: &[usize], k: usize) -> bool {
    TwoPointerStateIter::<SumOps>::new(a, (0, k)).any(|(_, _, (sum, _))| sum == k)
}

C. 有効な区間と内部状態の参照を返すイテレーター

fn find_sum_c(a: &[usize], k: usize) -> bool {
    let mut iter = TwoPointerLendingIter::<SumOps>::new(a, (0, k));
    while let Some((_, _, &(sum, _))) = iter.next() {
        if sum == k {
            return true;
        }
    }
    false
}

D. しゃくとり法ライブラリなし

fn find_sum_d(a: &[usize], k: usize) -> bool {
    let n = a.len();

    let mut l = 0;
    let mut s = 0;
    for r in 0..n {
        s += a[r];
        while s > k {
            s -= a[l];
            l += 1;
        }
        if s == k {
            return true;
        }
    }

    false
}

前の問題と同じように、有効な区間を探した後に、合計が k になるものが存在するかを、イテレーターの any() メソッドで調べます。

状態更新の例:

assert_eq!(find_sum(&[1, 2, 3, 4, 5], 10), true); // [1, 2, 3, 4].sum() = 10

[0]	[1]	[2]	[3]	state	l	r	check
				(0, 9)	0	0
1				(1, 9)	0	1	1 - 0 < 9
1	2			(3, 9)	0	2	3 - 0 < 9
1	2	3		(6, 9)	0	3	6 - 0 < 9
1	2	3	4	(10, 9)	0	4
	2	3	4	(9, 9)	1	4	10 - 1 = 9

「B. 有効な区間と内部状態のコピーを返すイテレーター」を使うと簡単に解けます。

「A. 有効な区間のみ返すイテレーター」の方法では、合計値が 9 になるかを区間から素早く判定できるように、イテレーターを呼ぶ側で準備が必要です。累積和を使います。 a[1..4].sum() は cum[4] - cum[1] と同じです。

	[0]	[1]	[2]	[3]	[4]	[5]
a	1	2	3	4	5
cum	0	1	3	6	10	15

ところで、「事前に累積和を計算する」と、「1. 単調増加な数列から 2 つの値を選び、差が k になるものがあるかを求める」と同じ問題になります。演算方法も同じものを使えます。

3. 数列の区間のうち、合計値が `10` 以下になる個数を求める

単体テスト:

#[test]
fn test_count_sum10_or_less() {
    assert_eq!(count_sum10_or_less(&[1, 2, 3, 4, 5]), 12);
    assert_eq!(count_sum10_or_less(&[1, 3, 3, 4, 5]), 11);
    assert_eq!(count_sum10_or_less(&[1, 3, 0, 4, 5]), 13);
    assert_eq!(count_sum10_or_less(&[1, 100, 1, 9, 1]), 6);
    assert_eq!(count_sum10_or_less(&[]), 0);
}

演算方法:

struct Sum10Ops;
impl TwoPointerIterOps for Sum10Ops {
    type I = usize;
    type S = usize;

    fn test(&state: &Self::S, _: &[Self::I]) -> bool {
        state <= 10
    }
    fn on_push_back(state: &mut Self::S, _: usize, &x: &Self::I) {
        *state += x;
    }
    fn on_pop_front(state: &mut Self::S, _: usize, &x: &Self::I) {
        *state -= x;
    }
}

状態: state == sum 区間の合計

A. 有効な区間のみ返すイテレーター:

fn count_sum10_or_less_a(a: &[usize]) -> usize {
    TwoPointerIter::<Sum10Ops>::from(a)
        .map(|(l, r)| r - l)
        .sum()
}

B. 有効な区間と内部状態のコピーを返すイテレーター:

fn count_sum10_or_less_b(a: &[usize]) -> usize {
    TwoPointerStateIter::<Sum10Ops>::from(a)
        .map(|(l, r, _)| r - l)
        .sum()
}

C. 有効な区間と内部状態の参照を返すイテレーター

fn count_sum10_or_less_c(a: &[usize]) -> usize {
    let mut count = 0;
    let mut iter = TwoPointerLendingIter::<Sum10Ops>::from(a);
    while let Some((l, r, _)) = iter.next() {
        count += r - l;
    }
    count
}

D. しゃくとり法ライブラリなし

fn count_sum10_or_less_d(a: &[usize]) -> usize {
    let n = a.len();
    let mut count = 0;

    let mut l = 0;
    let mut s = 0;
    for r in 0..n {
        s += a[r];
        while s > 10 {
            s -= a[l];
            l += 1;
        }
        count += r + 1 - l;
    }

    count
}

状態更新の例:

assert_eq!(count_sum10_or_less(&[1, 2, 3, 4, 5]), 12);

[0]	[1]	[2]	[3]	[4]	state	l	r	10以下の区間
1					1	0	1	[1]
1	2				3	0	2	[1, 2] [2]
1	2	3			6	0	3	[1, 2, 3] [2, 3] [3]
1	2	3	4		10	0	4	[1, 2, 3, 4] [2, 3, 4] [3, 4] [4]
1	2	3	4	5	15	0	5
	2	3	4	5	14	1	5
		3	4	5	12	2	5
			4	5	9	3	5	[4, 5] [5]

しゃくとり法でそれぞれの l に対して「合計が10以下になる最長の r」を求めます。「合計が10以下になる区間」は l < m <= r となる [l, m) で、その個数は r - l です。

イテレーターを終点まで回し、それぞれの r - l の合計値を求めます。これが「数列の区間のうち、合計値が 10 以下になる個数」となります。区間が 12個あると求まりました。

状態を取り出して使うことはありません。シンプルな「A. 有効な区間のみ返すイテレーター」で十分です。

この合計値の求め方は、「l に対して最大一回だけ行う派」「r に対して最大一回だけ行う派」のしゃくとり法向けです。「状態の判定を l が同じでも r を一つ動かすたびに行う派」の場合は同じ区間を二重カウントしないように、伸ばす方向がどちらかということを windows(2) で調べて分岐するような、ひと手間が増えます。

4. 数列の区間のうち、単調増加となる区間の個数を求める (ABC038-C)

C - 単調増加

単体テスト:

#[test]
fn test_abc038_c() {
    assert_eq!(abc038_c(&[1, 2, 3, 2, 1]), 8);
    assert_eq!(abc038_c(&[1, 2, 3, 4]), 10);
    assert_eq!(abc038_c(&[3, 3, 4, 1, 2, 2]), 8);
    assert_eq!(abc038_c(&[1, 5, 2, 3, 4, 2]), 10);

    let a: Vec<_> = (0..10000).collect();
    assert_eq!(abc038_c(&a), 10000 * 10001 / 2);
}

演算方法:

struct Abc038COps;
impl TwoPointerIterOps for Abc038COps {
    type I = usize;
    type S = ();

    fn test(_: &Self::S, a: &[Self::I]) -> bool {
        let n = a.len();
        n == 1 || a[n - 2] < a[n - 1]
    }
}

状態: () なし
- on_push_back(), on_pop_front() は状態を持たないため設定不要

A. 有効な区間のみ返すイテレーター:

fn abc038_c_a(a: &[usize]) -> usize {
    TwoPointerIter::<Abc038COps>::from(a)
        .map(|(l, r)| r - l)
        .sum()
}

B. 有効な区間と内部状態のコピーを返すイテレーター:

fn abc038_c_b(a: &[usize]) -> usize {
    TwoPointerStateIter::<Abc038COps>::from(a)
        .map(|(l, r, _)| r - l)
        .sum()
}

C. 有効な区間と内部状態の参照を返すイテレーター

fn abc038_c_c(a: &[usize]) -> usize {
    let mut count = 0;
    let mut iter = TwoPointerLendingIter::<Abc038COps>::new(a, ());
    while let Some((l, r, _)) = iter.next() {
        count += r - l;
    }
    count
}

D. しゃくとり法ライブラリなし

fn abc038_c_d(a: &[usize]) -> usize {
    let n = a.len();
    let mut count = 0;

    let mut l = 0;
    for r in 0..n {
        if l < r && a[r - 1] >= a[r] {
            l = r;
        }
        count += r + 1 - l;
    }

    count
}

状態更新の例:

assert_eq!(abc038_c(&[1, 5, 2, 3, 4, 2]), 10);

[0]	[1]	[2]	[3]	[4]	[5]	state	l	r	昇順の区間
1						()	0	1	[1]
1	5					()	0	2	[1, 5]
1	5	2				()	0	3
	5	2				()	1	3
		2				()	2	3	[2]
		2	3			()	2	4	[2, 3] [3]
		2	3	4		()	2	5	[2, 3, 4] [3, 4] [4]
		2	3	4	2	()	2	6
			3	4	2	()	3	6
				4	2	()	4	6
					2	()	5	6	[2]

区間の右 2 つが単調増加な関係になっている間、区間をのばし続けるようにします。これにより、しゃくとり法でそれぞれの l に対して「単調増加となる最長の r」がが求まります。「単調増加となる区間」は l < m <= r となる [l, m) で、その個数は r - l です。

イテレーターを終点まで回し、それぞれの r - l の合計値を求めます。これが「単調増加となる区間の個数」となります。区間が 10個あると求まりました。

状態を取り出して使うことはありません。シンプルな「A. 有効な区間のみ返すイテレーター」で十分です。

5. 数列の区間のうち、最大値と最小値の差が `k` 以下となる最大の長さを求める

単体テスト:

#[test]
fn test_find_max_range() {
    assert_eq!(find_max_range(&[1, 2, 3, 4, 5], 1), 2); // [1, 2]
    assert_eq!(find_max_range(&[1, 1, 2, 4, 2], 2), 3); // [1, 1, 2]
    assert_eq!(find_max_range(&[1, 2, 3, 4, 4, 3, 2, 1, 2, 3], 2), 6); // [2, 3, 4, 4, 3, 2]

    let a: Vec<_> = (0..20000).collect();
    assert_eq!(find_max_range(&a, 10000), 10001); // [0, 1, ..., 10000]
}

演算方法:

struct RangeOps;
impl TwoPointerIterOps for RangeOps {
    type I = usize;
    type S = (BTreeSet<(usize, usize)>, usize);

    fn test((set, k): &Self::S, _: &[Self::I]) -> bool {
        set.last().unwrap().0 - set.first().unwrap().0 <= *k
    }
    fn on_push_back((set, _): &mut Self::S, i: usize, &x: &Self::I) {
        set.insert((x, i));
    }
    fn on_pop_front((set, _): &mut Self::S, i: usize, &x: &Self::I) {
        set.remove(&(x, i));
    }
}

状態: (map, k)
- map: 区間に現れる数字と、その個数のマップ
- k: 区間の判定に使う固定値

A. 有効な区間のみ返すイテレーター:

fn find_max_range_a(a: &[usize], k: usize) -> usize {
    let a: Vec<_> = a.iter().enumerate().map(|(i, &x)| (x, i)).collect();
    TwoPointerIter::<RangeOps>::new(&a, (BTreeSet::new(), k))
        .map(|(l, r)| r - l)
        .max()
        .unwrap()
}

B. 有効な区間と内部状態のコピーを返すイテレーター (※重い):

fn find_max_range_b(a: &[usize], k: usize) -> usize {
    let a: Vec<_> = a.iter().enumerate().map(|(i, &x)| (x, i)).collect();
    TwoPointerStateIter::<RangeOps>::new(&a, (BTreeSet::new(), k))
        .map(|(l, r, _)| r - l)
        .max()
        .unwrap()
}

C. 有効な区間と内部状態の参照を返すイテレーター

fn find_max_range_c(a: &[usize], k: usize) -> usize {
    let a: Vec<_> = a.iter().enumerate().map(|(i, &x)| (x, i)).collect();
    let mut max_len = 0;
    let mut iter = TwoPointerLendingIter::<RangeOps>::new(&a, (BTreeSet::new(), k));
    while let Some((l, r, _)) = iter.next() {
        max_len = max_len.max(r - l);
    }
    max_len
}

D. しゃくとり法ライブラリなし

fn find_max_range_d_set(a: &[usize], k: usize) -> usize {
    let n = a.len();
    let mut max_len = 0;

    let mut l = 0;
    let mut s: BTreeSet<(usize, usize)> = BTreeSet::new(); // multiset
    for r in 0..n {
        s.insert((a[r], r));
        while s.last().unwrap().0 - s.first().unwrap().0 > k {
            s.remove(&(a[l], l));
            l += 1;
        }
        max_len = max_len.max(r + 1 - l);
    }

    max_len
}

assert_eq!(find_max_range(&[1, 2, 3, 4, 4, 3, 2, 1, 2, 3], 2), 6); // [2, 3, 4, 4, 3, 2]

[0]	[1]	[2]	[3]	[4]	[5]	[6]	[7]	[8]	[9]	state set	l	r	check
1										(1 ,0)	0	1	1
1	2									(1, 0) (2, 1)	0	2	2
1	2	3								(1, 0) (2, 1) (3, 2)	0	3	3
1	2	3	4							(1, 0) (2, 1) (3, 2) (4, 3)	0	4
	2	3	4							(2, 1) (3, 2) (4, 3)	1	4	3
	2	3	4	4						(2, 1) (3, 2) (4, 3) (4, 4)	1	5	4
	2	3	4	4	3					(2, 1) (3, 2) (3, 5) (4, 3) (4, 4)	1	6	5
	2	3	4	4	3	2				(2, 1) (2, 6) (3, 2) (3, 5) (4, 3) (4, 4)	1	7	6
	2	3	4	4	3	2	1			(1, 7) (2, 1) (2, 6) (3, 2) (3, 5) (4, 3) (4, 4)	1	8
		3	4	4	3	2	1			(1, 7) (2, 6) (3, 2) (3, 5) (4, 3) (4, 4)	2	8
			4	4	3	2	1			(1, 7) (2, 6) (3, 5) (4, 3) (4, 4)	3	8
				4	3	2	1			(1, 7) (2, 6) (3, 5) (4, 4)	4	8
					3	2	1			(1, 7) (2, 6) (3, 5)	5	8	3
					3	2	1	2		(1, 7) (2, 6) (2, 8) (3, 5)	5	9	4
					3	2	1	2	3	(1, 7) (2, 6) (2, 8) (3, 5) (3, 9)	5	10	5

状態の持ち方

これまでの 4問に比べると、状態の持ち方に工夫が必要です。

BTreeSet で現れる数字を管理します。そうすると、最大値と最小値の差が k 以下になるかは、 BTreeSet の末尾 (最大値) と先頭 (最小値) を引くことで素早く調べられます。

単純な BTreeSet では、同じ値が1つにまとめられます。区間 [3, 4, 3, 2, 1] の内部状態を { 1, 2, 3, 4 } のように持ってしまうと、左端の 3 を消したときに { 1, 2, 4 } のように途中の 3 までまとめて消してしまいかねません。これは嬉しくありません。

競技プログラミングでは、よく次のどちらかで対応します ⁴。今回は実装が簡単な上の方で対応します ⁵。

値とインデックスのペアを持つ set を用意する
- 区間 [3, 4, 3, 2, 1] は { (1, 0), (2, 3), (3, 0), (3, 2), (4, 1) } として持つ
- (3, 0) を取り除くと { (1, 0), (2, 3), (3, 2), (4, 1) } になる
- 次に (4, 1) を取り除くと { (1, 0), (2, 3), (3, 2) } になる
値と出現数の map を用意する
- 区間 [3, 4, 3, 2, 1] は { 1: 1, 2: 1, 3: 2, 4: 1 } として持つ
- (3, 2) を取り除くと、3 の出現数を一つ減らした { 1: 1, 2: 1, 3: 1, 4: 1 } になる
- 次に (4, 1) を取り除くと、4 が消えるのでキーからも除き、{ 1: 1, 2: 1, 3: 1 } になる

状態のコピー

「B. 有効な区間と内部状態のコピーを返すイテレーター」の紹介時に、マップのコピーが重いと書きました。これは BTreeSet についても同様です。テストケースにある 10000 要素を入れるような実装になると、手元では数秒かかりました。 ⁶

この問題を解く際に、状態を取り出して使うことはありません。シンプルな「A. 有効な区間のみ返すイテレーター」で十分です。

6. 数列の区間のうち、現れる数字が `k` 種類以下となる最大の長さを求める (典型90-034)

034 - There are few types of elements（★4）

単体テスト:

fn test_typical_034() {
    assert_eq!(typical_034(&[1, 2, 3, 4, 5], 1), 1); // [1]
    assert_eq!(typical_034(&[1, 1, 2, 4, 2], 4), 5); // [1, 1, 2, 4, 2]
    assert_eq!(typical_034(&[1, 2, 3, 4, 4, 3, 2, 1, 2, 3], 2), 4); // [3, 4, 4, 3]

    let a: Vec<_> = (0..20000).collect();
    assert_eq!(typical_034(&a, 10000), 10000); // [0, 1, ..., 9999]
}

演算方法:

struct Typical034Ops;
impl TwoPointerIterOps for Typical034Ops {
    type I = usize;
    type S = (BTreeMap<usize, usize>, usize);

    fn test((map, k): &Self::S, _: &[Self::I]) -> bool {
        map.len() <= *k
    }
    fn on_push_back((map, _): &mut Self::S, _: usize, x: &Self::I) {
        *map.entry(*x).or_insert(0) += 1;
    }
    fn on_pop_front((map, _): &mut Self::S, _: usize, x: &Self::I) {
        let Some(count) = map.get_mut(x) else { unreachable!() };
        *count -= 1;
        if *count == 0 {
            map.remove(&x);
        }
    }
}

状態: (map, k)
- map: 区間に現れる数字と、その個数のマップ
- k: 区間の判定に使う固定値

A. 有効な区間のみ返すイテレーター:

fn typical_034_a(a: &[usize], k: usize) -> usize {
    TwoPointerIter::<Typical034Ops>::new(a, (BTreeMap::new(), k))
        .map(|(l, r)| r - l)
        .max()
        .unwrap()
}

B. 有効な区間と内部状態のコピーを返すイテレーター (※重い):

fn typical_034_b(a: &[usize], k: usize) -> usize {
    TwoPointerStateIter::<Typical034Ops>::new(a, (BTreeMap::new(), k))
        .map(|(l, r, _)| r - l)
        .max()
        .unwrap()
}

C. 有効な区間と内部状態の参照を返すイテレーター

fn typical_034_c(a: &[usize], k: usize) -> usize {
    let mut max_len = 0;
    let mut iter = TwoPointerLendingIter::<Typical034Ops>::new(a, (BTreeMap::new(), k));
    while let Some((l, r, _)) = iter.next() {
        max_len = max_len.max(r - l);
    }
    max_len
}

D. しゃくとり法ライブラリなし

fn typical_034_d(a: &[usize], k: usize) -> usize {
    let n = a.len();
    let mut max_len = 0;

    let mut l = 0;
    let mut s: BTreeMap<usize, usize> = BTreeMap::new();
    for r in 0..n {
        *s.entry(a[r]).or_insert(0) += 1;
        while s.len() > k {
            let x = a[l];
            match s.get(&x).unwrap() {
                &1 => s.remove(&x),
                &m => s.insert(x, m - 1),
            };
            l += 1;
        }
        max_len = max_len.max(r - l + 1);
    }

    max_len
}

状態更新の例:

assert_eq!(typical_034(&[1, 2, 3, 4, 4, 3, 2, 1, 2, 3], 2), 4); // [3, 4, 4, 3]

[0]	[1]	[2]	[3]	[4]	[5]	[6]	[7]	[8]	[9]	state map	l	r	check
1										{1: 1}	0	1	1
1	2									{1: 1, 2: 1}	0	2	2
1	2	3								{1: 1, 2: 1, 3: 1}	0	3
	2	3								{2: 1, 3: 1}	1	3	2
	2	3	4							{2: 1, 3: 1, 4: 1}	1	4
		3	4							{3: 1, 4: 1}	2	4	2
		3	4	4						{3: 1, 4: 2}	2	5	3
		3	4	4	3					{3: 2, 4: 2}	2	6	4
		3	4	4	3	2				{2: 1, 3: 2, 4: 2}	2	7
			4	4	3	2				{2: 1, 3: 1, 4: 2}	3	7
				4	3	2				{2: 1, 3: 1, 4: 1}	4	7
					3	2				{2: 1, 3: 1}	5	7	2
					3	2	1			{1: 1, 2: 1, 3: 1}	5	8
						2	1			{1: 1, 2: 1}	6	8	2
						2	1	2		{1: 1, 2: 2}	6	9	3
						2	1	2	3	{1: 1, 2: 2, 3: 1}	6	10
							1	2	3	{1: 1, 2: 1, 3: 1}	7	10
								2	3	{2: 1, 3: 1}	8	10	2

「5. 数列の区間のうち、最大値と最小値の差が k 以下となる最大の長さを求める」で状態を扱う 2つの方法を紹介しました。今回の問題では種類の個数を知りたいですので、出現する数とその個数の map を使います。

例えば map[2]=1 は、区間内に数字 2 が 1つ含まれることを示します。そして、現れる数の種類は map.len() で求められます。map[2]=0 のように個数 0 のものをカウントし続けないように、0 になるときはキーから除くようにしています。

これまでと同じように演算方法を設定することで、6. も解けました。

最後に

長文記事にお付き合いいただきありがとうございました。

本記事では以下を行いました:

しゃくとり法の簡単な紹介と、複数の関数に分ける考え方を示します
Rust のイテレーターで扱うライブラリを 3通り実装します
- 「A. 有効な区間のみ返す」Iterator トレイトを満たす実装
- 「B. 有効な区間と内部状態のコピーを返す」 Iterator トレイトを満たす実装
- 「C. 有効な区間と内部状態の参照を返す」 Rust 1.65 GATs の LendingIterator 例を満たす実装
しゃくとり法の例題 6問を、A, B, C と「D. しゃくとり法ライブラリなし」の 4通りで解きます
1. 単調増加な数列から 2 つの値を選び、差が k になるものがあるかを求める
2. 数列の区間のうち、合計値が k になるものがあるかを求める
3. 数列の区間のうち、合計値が 10 以下になる個数を求める
4. 数列の区間のうち、単調増加となる区間の個数を求める (ABC038-C)
5. 数列の区間のうち、最大値と最小値の差が k 以下となる最大の長さを求める
6. 数列の区間のうち、現れる数字が k 種類以下となる最大の長さを求める (典型90-034)

しゃくとり法をライブラリ化すると便利かもと思った方は、使ってみる、好みの実装を書いてみることをお勧めします。「A. 有効な区間のみ返す」だけでも用意すると、楽になると思います。

Rust では所有権の面倒さがありました。C++ など他の言語では、より簡単に汎用性の高いライブラリを書けると思います。

まあでも、考え方が整理されれば、ライブラリなしにべたっと書くのもそんなに難しくない、というところですかと。

このライブラリが実用的かどうかは分かりません。「しゃくとり法を部品に分けて考える」ことの練習にはなると思います。そんな紹介記事、兼、2024年エイプリルフールのネタ記事でした。

おまけ: しゃくとり法以外の話

本記事では Two-Pointer でしゃくとり法を行いました。他にも Two-Pointer 2つの位置の動かし方は、しゃくとり法に限りません。

ここからはしゃくとり法以外の話です。

7. 数列から 2 つの値を選び、積が `k` 以下になる組の個数を求める

「数列から 2 つの値を選び、積が k 以下になる組の個数を求める」問題を考えます。

単体テスト:

#[test]
fn test_count_mul() {
    assert_eq!(count_mul(&[3, 1, 4, 1, 5], 12), 8);
    assert_eq!(count_mul(&[3, 1, 4, 1, 5], 15), 9);
    assert_eq!(count_mul(&[1, 2, 2, 3, 3, 4], 6), 10);
    assert_eq!(count_mul(&[2, 3, 4, 5, 6], 5), 0);
    assert_eq!(count_mul(&[], 5), 0);

    let a: Vec<_> = (0..10000).collect();
    assert_eq!(count_mul(&a, 1), 9999); // (0, 1), ..., (0, 9999)
    assert_eq!(count_mul(&a, 2), 10000); // (0, 1), ..., (0, 9999), (1, 2)
    assert_eq!(count_mul(&a, 9999 * 9998), 10000 * 9999 / 2); // all
    assert_eq!(count_mul(&a, 9998 * 9998), 10000 * 9999 / 2 - 1); // all except (9998, 9999)
}

両端の積が k 以下になるように縮めていき、長さが 0 になったところで終了です。組の数え方は「3. 数列の区間のうち、合計値が 10 以下になる個数を求める」同様です。

状態更新の例:

assert_eq!(count_mul(&[1, 2, 2, 3, 3, 4], 6), 10);

[0]	[1]	[2]	[3]	[4]	[5]	l	r	check
1	2	2	3	3	4	6	0	(1, 2) (1, 2) (1, 3) (1, 3) (1, 4)
	2	2	3	3		1	5	(2, 2) (2, 3) (2, 3)
		2	3	3		2	5	(2, 3) (2, 3)

このような右の範囲を縮めるタイプの問題については、しゃくとり法ライブラリは対象外です。対象外ですが、ここまで記事にお付き合いいただいた方は、似たようなものだと考えて実装できるでしょう。

Two-pointer での実装例

fn count_mul(a: &[usize], k: usize) -> usize {
    let a: Vec<_> = a.iter().sorted().copied().collect();
    let n = a.len();

    let mut r = n;
    let mut count = 0usize;
    for l in 0..n {
        while l < r && a[l] * a[r - 1] > k {
            r -= 1;
        }
        if l == r {
            break;
        }
        count += r - l - 1;
    }

    count
}

各問題のしゃくとり法を使わない別解

今回取り上げた問題 7問のうち 5問は、しゃくとり法を使わなくても短い実装で解けます。

1. 単調増加な数列から 2 つの値を選び、差が `k` になるものがあるかを求める - 二分探索

fn find_sub_e(a: &[usize], k: usize) -> bool {
    a.iter()
        .enumerate()
        .any(|(i, &x)| a[(i + 1)..].binary_search(&(x + k)).is_ok())
}

2. 数列の区間のうち、合計値が `k` になるものがあるかを求める - 累積和と二分探索

fn find_sum_e(a: &[usize], k: usize) -> bool {
    let mut cum = vec![0; a.len() + 1];
    for (i, &x) in a.iter().enumerate() {
        cum[i + 1] = cum[i] + x;
    }
    cum.iter()
        .enumerate()
        .any(|(i, &x)| cum[(i + 1)..].binary_search(&(x + k)).is_ok())
}

3. 数列の区間のうち、合計値が `10` 以下になる個数を求める - 累積和と二分探索

fn count_sum10_or_less_e(a: &[usize]) -> usize {
    let mut cum = vec![0; a.len() + 1];
    for (i, &x) in a.iter().enumerate() {
        cum[i + 1] = cum[i] + x;
    }
    cum.iter()
        .enumerate()
        .map(|(i, &x)| cum[(i + 1)..].partition_point(|&y| y - x <= 10))
        .sum()
}

4. 数列の区間のうち、単調増加となる区間の個数を求める (ABC038-C) - 1重ループ

fn abc038_c_e(a: &[usize]) -> usize {
    let mut v = vec![1usize];
    for a in a.windows(2) {
        if a[0] < a[1] {
            *v.last_mut().unwrap() += 1;
        } else {
            v.push(1);
        }
    }
    v.iter().map(|&x| x * (x + 1) / 2).sum()
}

5. 数列の区間のうち、最大値と最小値の差が `k` 以下となる最大の長さを求める - セグメント木

use ac_library::{Monoid, Segtree};

pub struct MinMax(Infallible);
impl Monoid for MinMax {
    type S = (usize, usize);
    fn identity() -> Self::S {
        (usize::MAX, usize::MIN)
    }
    fn binary_operation(&(a0, a1): &Self::S, &(b0, b1): &Self::S) -> Self::S {
        (a0.min(b0), a1.max(b1))
    }
}

fn find_max_range_e(a: &[usize], k: usize) -> usize {
    let n = a.len();
    let v: Vec<_> = a.iter().map(|&x| (x, x)).collect();
    let segtree = Segtree::<MinMax>::from(v);
    (0..n)
        .map(|l| segtree.max_right(l, |&(min, max)| max < min || max - min <= k) - l)
        .max()
        .unwrap()
}

7. 数列から 2 つの値を選び、積が `k` 以下になる組の個数を求める - 二分探索

fn count_mul_e(a: &[usize], k: usize) -> usize {
    let a: Vec<_> = a.iter().sorted().copied().collect();

    let mut count = 0usize;
    for (l, &x) in a.iter().enumerate() {
        let len = a[(l + 1)..].partition_point(|&y| x * y <= k);
        if len == 0 {
            break;
        }
        count += len;
    }

    count
}

区間全体から内部状態を求める 5., 6. では、しゃくとり法が有効です。 5. はセグメント木を使って任意区間の最大最小値を求められるようにした例を載せましたが、これはしゃくとり法の方が簡単だと思います。

逆に言うと、内部状態が単純なものについては、しゃくとり法は実装の手間に見合わないかもしれません。 $O(N log N)$ の二分探索より $O(N)$ のしゃくとり法の方が少し高速とは言えます。制限時間内に解けるかどうかというと誤差の範囲です。とくに、内部状態が複雑でもうまく値を参照できるような LendingIterator の出番はいったい。

そんな、エイプリルフールのネタらしい記事でした。

本記事の実装は、3種類のイテレーターで実装を使いまわせる、という基準で選んでいます。普段書くときにこの方法がおすすめ、というものではありません。 ↩
lending_iterator - Rust のように、any() などの Iterator 機能の一部を使えるクレートもあります。将来的には std でこのようなことができると嬉しいです。 ↩
2つの値の差という問題なのに、このように区間の長さ 1 を評価していることに違和感を持たれる方もいると思います。これは問題ありません。区間の長さが 1 ということは、先頭と終端の 2点が同じ場所を指しています。今回の問題では、同じ場所を指したときの差は 0 です。制約 k > 0 ならこちらで十分です。気になる場合、r - l > 1 という条件文を書くのも良いと思います。 ↩
Rust の標準ライブラリには C++ の multiset がないことと、multiset 相当を使うとしてもキー単位で重複要素がまとめて消えるとこの問題では困る、という理由です ↩
下の方針にすると、本記事中で on_push_back(), on_pop_front() でインデックスを使う機会がなくなる、という理由もあります。 ↩
コピーした結果を (l, r, _) と呼び飛ばしていれば、もしかすると最適化でコピーが省略されるかも…… と期待しました。手元の Rust 1.70 では、そんな都合の良い動作にはなりませんでした。 ↩