アルゴリズム - 整数上の二分探索

概要

整数（のようなデータ）について Yes/No を返す述語 \(P\) があるとする: \[P \colon \mathbb Z \to \mathrm{Bool}.\]

そして今, この \(P\) はある整数 \(m\) があって,

\(n < m \implies P(n) = \mathrm{No}\)
\(n \geq m \implies P(n) = \mathrm{Yes}\)

を満たすとする. このとき, この整数 \(m\) を求めたい.

ただし, 次のような2つの値 \(l, r (l < r)\) が予め与えられるとする:

\(P(l) = \mathrm{No}\),
\(P(r) = \mathrm{Yes}\).

解法

区間 \((l,r]\) は \(m\) を含んでいる. このことを不変条件に持つように上手く区間のサイズを半分にしてく. そのためには \(l, r\) の適当な中間値を持ってきて, それが \(P\) を満たすかをチェックするだけでいい. これを繰り返して, 区間のサイズがちょうど \(1\) になったとき, その要素が求める答え \(m\) である.

ところで \(l,r,m\) の乗ってるデータは, 中間値を取る操作 middle と, 区間のサイズが \(1\) であることをチェックする操作 close を必要とする. 逆に言えばこの二つさえあれば整数そのものに限らなくて良い. 例えば十分小さい値 eps を定めて close(l, r) = (r - l < eps) とすることで浮動小数点数であっても, 精度 eps で m が求まる.

応用

整数として配列のインデックス (usize) を選び, prop を上手く作ることで, 昇順ソート済みの配列 xs 中に x がいくつあるか, x 以上/以下がいくつあるか, などを対数時間で計算できる.

/// Algorithm - Binary Search (lowerbound)
pub trait CompleteIdx: Copy {
    fn mid(self, other: Self) -> Self;
}
#[macro_export]
macro_rules! completeidx {
    ( $type:ty, mid($self:ident, $other:ident) = $code:block ) => {
        impl CompleteIdx for $type {
            fn mid($self, $other: Self) -> Self { $code }
        }
    };
}
completeidx! { usize, mid(self, other) = { (self + other) / 2 }}
completeidx! { u128, mid(self, other) = { (self + other) / 2 }}
completeidx! { i128, mid(self, other) = { (self + other) / 2 }}
completeidx! { u64, mid(self, other) = { (self + other) / 2 }}
completeidx! { i64, mid(self, other) = { (self + other) / 2 }}
completeidx! { f64, mid(self, other) = { (self + other) / 2.0 }}

pub fn lowerbound<T: CompleteIdx>(r: std::ops::Range<T>, cond: &dyn Fn(T) -> bool) -> Option<T> {
    if cond(r.start) {
        return Some(r.start);
    }
    // TODO(from 1.47.0)
    // if r.is_empty() { return None }
    let mut left = r.start;
    let mut right = r.end;
    let mut ok = false;
    for _ in 0..100 {
        let mid = T::mid(left, right);
        if cond(mid) {
            right = mid;
            ok = true;
        } else {
            left = mid;
        }
    }
    if ok {
        Some(right)
    } else {
        None
    }
}

#[cfg(test)]
mod test_binary_search {
    use crate::algorithm::binary_search::*;

    #[test]
    fn search_bound() {
        let v: Vec<i32> = (0..100).collect();
        assert_eq!(lowerbound(0..100, &|i| v[i] > 50), Some(51));
        assert_eq!(lowerbound(0..100, &|i| v[i] >= 0), Some(0));
        assert_eq!(lowerbound(0..100, &|i| v[i] >= 99), Some(99));
        assert_eq!(lowerbound(0..100, &|i| v[i] >= 100), None);
    }

    #[test]
    fn search_on_real_number() {
        let x: f64 = lowerbound(0.0..2.0, &|x| x * x >= 2.0).unwrap();
        assert!((x * x - 2.0).abs() < 0.00001);
    }
}