区間ヒープってバグりますよね。わかります。概略は簡単なのですけれども長さを $4$ で割った余りで謎の場合分けが入ってつらい印象です。この記事では私が簡単かつ高速であると思っている実装をご紹介していきます。

データ表現

この辺りは基本的にご存じな前提でお話を進めたいのですが、認識をそろえるために簡単にご説明します。

区間ヒープは根が最大値・最小値を持ち、残りを左の子・右の子で任意に半々（奇数ならば左が多い）に分割して管理するヒープ型木データ表現です。つまり各節点が基本的に $2$ 個（要素数が奇数のときは末尾だけ $1$ 個）の要素を管理する二分木の形をしていて、かつその節点の管理する $1$ ～ $2$ 個の要素の張る区間を考えると、包含関係に関してヒープ条件を満たすようなデータ表現とみることもできます。

さらにこれは配列 $a$ で表現することができます。各節点が $1$ ～ $2$ 要素分のスパンを管理するように二分ヒープ同様に配置すればよいです。このとき

［min-heap 条件］添字が偶数のものだけ見たときに min-heap になっている
［添字が例外的でない max-heap 条件］添字が奇数のものだけ見たときに max-heap になっている
［添字が例外的な max-heap 条件］$n$ が $3$ 以上の奇数ならば $a _ { n - 1 }$ はその親節点の右側要素 $a _ { 2 \left \lfloor \frac { n - 2 } { 4 } \right \rfloor + 1 }$ 以下である
［区間条件］$a _ { 2i } \le a _ { 2i + 1 }$

が区間ヒープ条件です。この「添え字が例外的な max-heap 条件」というのが曲者で、雑に実装するとこれが壊れがちです。

最小値削除・最大値削除

最小値削除（末尾周辺の面倒な場合分け以外）

二分ヒープ同様最小値削除クエリでは先頭と末尾とを入れ替え末尾を削除し、そしてこれにより最小でなくなった先頭をどんどん葉の方向に再帰的に動かしていきます。

再帰の各ステップの詳細を、末尾周辺の面倒な場合分け以外、すなわち左右の子が存在してさらにそれらの最小値・最大値のフィールドが非空なときを考えましょう。すなわち次の図のような場合です。

現在の節点の管理する区間が $[m, M]$、左右の子の管理する区間が $[m _ L, M _ L], [m _ r, M _ R]$ であって「$m$ だけがおかしい」、つまり

左右の部分木は区間ヒープの条件を満たしてる
$\max ( M _ L, M _ R ) ≤ M$

である状況です。この状況から区間ヒープ条件を回復する方法を考えましょう。

$6$ つの値の大小関係は例えば次のようになっています。$[m _ L, M _ L], [m _ r, M _ R]$ はそれぞれ大小関係の正しい区間表現になっていて、かつ各々の子孫もその中に納まっており、$M$ はそれらすべて以上であるという状況ですが、$m$ の値だけが、どこにあるかわからない状況です。図では最悪の場合ということで、$m$ が $M$ よりも大きい場合を描いています。

ここからでも入れる保険をご紹介しましょう。まず $m > M$ ならばこれらを入れ替えることで $m ≤ M$ としてよいです。さらに $m > \min (m _ L, m _ R)$ ならばこれらの小さい方と入れ替えることで、$m$ が最小値になります。しかしそうすると今度は入れ替えに選ばれた方が最小値性を満たすとは限らなくなりますね。というわけでそちらに再帰すればよいです。

たとえば先程の図の場合は、次の図のように $R$ 側が壊れますので、そちらに再帰しましょう。

末尾周辺の面倒な場合分け

前のセクションでは $M, m _ L, M _ L, m _ R, M _ R$ がすべて揃っている前提でお話ししましたが、ここではそれが欠けている場合を含めて議論しましょう。先ほどのアルゴリズムを（しれっと面倒な場合にも対応できるように直して）厳密に書くと次のようになります。

$M$ が null なら終了
$m > M$ ならば $m$ と $M$ を入れ替える
$m _ L, m _ R$ のうち、non-null でかつ $m$ より小さいものが
1. あればその中で最小のもの $m _ X (X = L, R)$ をとる
2. なければ終了
$X$ に再帰する

これの正当性を示しましょう。第 2 行によって追加で $m < M$ が成り立ち、さらにこれにより $M$ が正しい最大値になることも保証されます。そして第 3 行によって $m$ も正しい最小値になりますが、代わりに $m _ X$ が正しい最小値ではなくなる可能性があります。しかしそれは第 4 行で修正されます。

最小値削除の配列表現

$m$ の入っている添字を $i$（コード中の start）とおくと、$M, m _ L, m _ R$ の添字はそれぞれ $i + 1, 2 i + 2, 2 i + 4$ と書けますから、次のように実装できます。

fn pop_min<T: Ord>(values: &mut Vec<T>) {
    (!values.is_empty()).then_some(())?;
    let ret = values.swap_remove(0);
    min_heapfy_down(values, 0);
    Some(ret)
}
fn min_heapify_down<T: Ord>(values: &mut [T], mut start: usize) {
    let n = values.len();
    loop {
        let end = start + 1;
        if end >= n { // 1
            break;
        }
        if values[start] > value[end] {
            values.swap(start, end); // 2
        }
        let mut next = 2 * start + 4;
        if next >= n || values[next - 2] < values[next ] {
            swp -= 2;
        }
        if next >= n || values[start] <= values[next] {
            break; // 3b
        }
        values.swap(start, next); // 3a
        start = swp; // 4
    }
}

最大値削除の場合の面倒なこと

実は今までしれっと要素が $1$ 個しかないノードにおいてはその要素が $m$ であるとみなしていたため、最小値削除はうまくいっていたのですが、最大値削除の場合にはそれは $M$ であるとみなさないと今までの議論は通らなくなります。とはいえ実装上そこに柔軟性を持たせるのは困難である（特に配列表現にできない）ため、いかなる場合も singleton は $m$ であるとみなすことにします。

すると再帰ステップが singleton の親になったときに最大値が片方無視されることになって、壊れてしまいます。例えば次のようなケースで一切スワップが起きず、区間ヒープ条件が回復しません。

そこで、これを各再帰ステップで真面目にハンドルしてもよいのですが、先述の通りこれで壊れるのは singleton とその親の関係だけなので、実は普通にやったあとそこだけ直せばよかったりして、こちらのほうがボトルネック部分が簡単になって嬉しいです。壊れていた場合は swap すれば治って、かつほかのところも壊れません。

fn pop_max<T: Ord>(values: &mut Vec<T>) -> Option<T> {
    if values.len() <= 2 {
        return self.values.pop();
    }
    let ret = values.swap_remove(1);
    max_heapify_down(&mut values, 1);
    Some(ret)
}

fn max_heapify_down<T: Ord>(values: &mut [T], mut end: usize) {
    let n = values.len();
    loop {
        let start = end - 1;
        if values[start] > values[end] {
            values.swap(start, end);
        }
        let mut next = 2 * end + 3;
        if next >= n || values[next - 2] > values[next] {
            next -= 2;
        }
        if next >= n || values[end] >= values[next] {
            break;
        }
        values.swap(end, next);
        end = next;
    }
    if n % 2 == 1 && 1 < n && values[n - 1] > values[(n / 2 - 1) | 1] {
        values.swap(n - 1, (n / 2 - 1) | 1); // singleton の親だけ直す
    }
}

ヒープの構築

配列が与えられたとき、線形時間で二分ヒープ条件を満たすようにできることはよく知られていると思います。区間ヒープでも深いことは考えずに先ほど定義した {min,max}_heapify_down を下から順に呼んでいけばよいのですが、正当性のためにはまだ一つだけ証明していないことがあります。余談ですが地味に二分ヒープのときと違ってコード中の (0..n).rev() の n を半分くらいにしようとするとバグります。なぜならば葉であっても先ほどの疑似コードの $m < M$ を解決する部分が必要だからです。

fn build_heap<T: Ord>(mut values: Vec<T>) {
    for i in (0..values.len()).rev() {
        match i % 2 {
            0 => min_heapify_down(&mut values, i),
            1 => max_heapify_down(&mut values, i),
            _ => unreachable!(),
        }
    }
}

二分木表現で考えて、あるノードを {min,max}_heapify_down するときには最小値削除のとき同様、部分木は区間ヒープ条件を満たしています。しかし今回は $m$ と $M$ の片方だけおかしいのを直すのではなくて、両方おかしい状態です。この状態からでも max_heapify_down と min_heapify_down を順番に呼べば治る」ことを示す必要があります。……といっても別に難しくはなくて先ほどの証明がこの場合でも通ることを見直せばよいです。各自ぜひです。

要素の挿入

二分ヒープのときと同様、とりあえず末尾に挿入してからそれを上に動かしていくことで条件を回復していきましょう。（力尽きて簡単な説明にとどめていますけれども実際丁寧に証明すると最小値・最大値削除と同じくらい面倒なはず）

挿入後に要素数が偶数になった場合　$m$ と $M$ を適宜入れ替えて $m \le M$ になるようにしましょう。すると、入れ替えなかった場合は $M$ が max-heap 条件から外れていて、入れ替えた場合は $m$ が min-heap 条件から外れているのでこれを二分ヒープのアルゴリズム（コード中の {min,max}_heapify_up）で修正しましょう。これにより $m < M$ などの区間条件も壊れないことがわかります。

挿入後に要素数が奇数になった場合　$m$ は min-heap 条件に違反しているかもしれませんし、（末尾なので添字が例外的な）max-heap 条件に違反しているかもしれません。親の $M$ の添字 $2 \left \lfloor \frac { n - 2 } { 4 } \right \rfloor + 1$ はビット演算で (n / 2 - 1) | 1 と書けるので実際に比較して後者であるかどうかを判定し、後者であれば添字が例外的な max-heap 条件だけ回復しておいて、残りは max_heapify_up をすればよいです。さもなくば前者であるかそもそも壊れていないかなので、min_heapify_up をすればよいです。

fn push<T: Ord>(values: &mut [T], x: T) {
    values.push(x);
    let n = values.len();
    match n % 2 {
        0 => {
            if values[n - 2] > values[n - 1] {
                values.swap(n - 2, n - 1);
                min_heapify_up(&mut values, n - 2);
            } else {
                max_heapify_up(&mut values, n - 1);
            }
        }
        1 => {
            if n == 1 {
                return;
            }
            let end = (n / 2 - 1) | 1;
            if values[end] < values[n - 1] {
                values.swap(end, n - 1);
                max_heapify_up(&mut values, end);
            } else {
                min_heapify_up(&mut values, n - 1);
            }
        }
        _ => unreachable!(),
    };
}