ブログ名

競技プログラミングやお歌のお話をする高菜です。

トップ > Convex hull trick などの実装に応用のある、Rust の BTreeSet / BTreeMap １本だけで、２通り以上の二分探索を実現する方法

Convex hull trick などの実装に応用のある、Rust の BTreeSet / BTreeMap １本だけで、２通り以上の二分探索を実現する方法

イントロ

Convex hull trick では直線の列を、傾きに関して単調増加になるように管理するのですが、二分探索の検索クエリは傾きだけでは足りません。次のように２種類の二分探索が欲しくなります。

  • 傾きに関して二分探索
  • 最小を達成する区間に関して二分探索

しかしあいにく私はこれを上手にやる方法を知らず、HashMap ２本とかで頑張るしかないかなと思っていました。ただ整合性を担保しないといけないのでバグりやすいですし、（今回は要素が Copy トレイトを実装しているのでそこまで気にしなくてもよいとはいえ）オブジェクトがコピーされるのは Rust 的には気持ちよくはないがちですよね。

そこで私はそれを Rust の力でなんとか１本にねじ込む方法を見つけましたので、本記事でご紹介しようと思います。（C++ などの他言語でできるのかわからないのでわかる方教えていただけるととても嬉しいです。）

問題

A, B を全順序集合、 C を集合とします。 列 (a _ i, b _ i, c _ i) _ { 0 \le i \lt n } ((a _ i, b _ i , c _ i) \in A \times B \times C) であって a _ i, b _ i の両方について狭義単調増加（i.e. a _ 0 \lt \dots \lt a _ { n - 1 }, b _ 0 \lt \dots \lt b _ { n - 1 }）なものに対して、次の操作を O ( \lg n )で行えるようなデータ構造を設計してください。

  1. x \in A が与えられるので、 x \le a _ i なる最小の i に対応する (a _ i, b _ i, c _ i) の取得や削除
  2. y \in B が与えられるので、 y \le b _ i なる最小の i に対応する (a _ i, b _ i, c _ i) の取得や削除
  3. 挿入可能な新しい要素 (a _ i, b _ i, c _ i) の挿入

要するに insert, delete, search がどちらのキーでもできてくださいということです。merge, split もやりたかったのですが、BTreeSet::append がおそらく（今回のように単に後ろに入れるだけの場合も）線形かかるのでやめておきました。むーんログマージほしいです……（単に後ろに入れるだけでないときの挙動どうするのですか問題はあるので用意しない気持ちもわかるのですが。）

やりかた

３つご紹介します。

1. ハックしない解法

a _ i \mapsto ( a _ i, b _ i, c _ i ) と b _ i \mapsto ( a _ i, b _ i, c _ i ) とで平衡２分木を２つ持ってばよいです。 難しくはないのですが、いちいち２回検索したり２回 mutate したりと、バグの温床になりやすいです。

2. 平衡二分木自作

当然できますね。 むしろ最近までは自作しないと無理かと思っていました。

3. 標準ライブラリの BTreeMap を使う

これで気づいた私天才かもしれません。たとえば、BTreeSet::get のシグニチャをご覧いただきたいのですが、このように検索クエリ Q はジェネリック引数です。ハックできます…！！

pub fn get<Q>(&self, value: &Q) -> Option<&T>
where
    T: Borrow<Q> + Ord,
    Q: Ord + ?Sized,

なお他にも BTreeSet::split_off, BTreeSet::remove, BTreeSet::range など様々なメソッドがこのインターフェースを採用しています。

というわけでタプル A \times B を意味する型 Pair を定義して、その中に A を意味する First、 B を意味する Second を格納しましょう。そして仕上げに Pair に Borrow<First> と Borrow<Second> を実装すれば完成です。BTreeMap に入れたいので Ord も実装しておきましょう。

#[derive(PartialEq, Eq, PartialOrd, Ord)]
struct Pair(First, Second);
#[derive(PartialEq, Eq, PartialOrd, Ord)]
struct First(i32);
#[derive(PartialEq, Eq, PartialOrd, Ord)]
struct Second(i32);
impl Borrow<First> for Pair {
    fn borrow(&self) -> &First {
        &self.0
    }
}
impl Borrow<Second> for Pair {
    fn borrow(&self) -> &Second {
        &self.1
    }
}

するとあら不思議。このようにどちらでも検索できます。

let mut map = vec![
    (Pair(First(1), Second(11)), 'a'),
    (Pair(First(3), Second(13)), 'b'),
    (Pair(First(5), Second(15)), 'c'),
]
.into_iter()
.collect::<BTreeMap<_, _>>();

// get
assert_eq!(map.get(&First(0)), None);
assert_eq!(map.get(&First(1)), Some(&'a'));
assert_eq!(map.get(&First(2)), None);
assert_eq!(map.get(&First(3)), Some(&'b'));
assert_eq!(map.get(&First(4)), None);
assert_eq!(map.get(&First(5)), Some(&'c'));
assert_eq!(map.get(&First(6)), None);

assert_eq!(map.get(&Second(10)), None);
assert_eq!(map.get(&Second(11)), Some(&'a'));
assert_eq!(map.get(&Second(12)), None);
assert_eq!(map.get(&Second(13)), Some(&'b'));
assert_eq!(map.get(&Second(14)), None);
assert_eq!(map.get(&Second(15)), Some(&'c'));
assert_eq!(map.get(&Second(16)), None);

// split_off
let mut other = map.split_off(&First(2));
assert_eq!(map.values().copied().collect::<String>().as_str(), "a");
assert_eq!(other.values().copied().collect::<String>().as_str(), "bc");
map.append(&mut other);

let mut other = map.split_off(&Second(15));
assert_eq!(map.values().copied().collect::<String>().as_str(), "ab");
assert_eq!(other.values().copied().collect::<String>().as_str(), "c");
map.append(&mut other);

// insert
map.insert(Pair(First(4), Second(14)), 'z');
assert_eq!(map.values().copied().collect::<String>().as_str(), "abzc");

この方式でできないこと

一応逆に自作すればできるけれどこの方式でできないこととして、 a _ i + b _ i をキーとした二分探索があります。Borrow::borrow はライフタイムの制約的にその場で計算した値への参照を返すことができないのです。必要ならタプル ` (a _ i, b _ i, a _ i + b _ i) の列を管理するしかないでしょうかね。

まとめ

検索クエリジェネリック引数、Vec<T> を [T] で検索したり String を str で検索したりするためくらいにしか思っていなかったのですが、案外使いみちが広かったことに気付かされました。

ところで Borrow ってこんな使い方をして良いのでしょうかね。Deref アンチパターン（Deref should only be implemented for smart pointers）のような感じでむやみに実装すると怒られの発生するものだったりしますでしょうかね。いずれにしても、インターフェースでは隠蔽して内部実装でやる分には、他に方法がないならよいのではないかとい気持ちになっています。

追記［2021-12-16］

C++ はできるそうです。shino さんありがとうございます。