その日 BinaryHeap は思い出しました。std に支配されていた恐怖を…… Reversed の中に囚われていた屈辱を！

みなさまは BinaryHeap<Reversed<i32>> がお嫌いです。

なぜかというと、いやまあ良いのですよ、型定義が複雑になる分には……

let mut heap = BinaryHeap<Reversed<i32>>::new();

heap.push(Reversed(42));
heap.push(Reversed(43));

let Reversed(x) = heap.pop().unwrap();
assert_eq!(x, 43);

なんですかこれ！？ 式中に Reversed とかいう 𝒂𝒃𝒔𝒕𝒓𝒂𝒄𝒕 𝒏𝒐𝒏𝒔𝒆𝒏𝒔𝒆 がコードにモリモリ登場するせいで、アルゴリズムが見づらくなってしまいます。

とくに min-heap と max-heap を両方保持してあっちへこっちへ移し替えるようなコードを書こうと思うと、本質ではないところが複雑になって、バグ埋め小宮さんです。

👇️ これなどです (ヒープの練習問題)

E - NarrowRectangles

というわけで、二分ヒープはご自身で実装しましょう。（完）

いやいや、でも標準ライブラリって速いんでしょう？ 自分で実装なんてしたら……

そんなことはありません！

確かに std の実装を読むと、何やら謎の工夫が散りばめられてて、とても太刀打ちできそうになりように見えます。

しかし、なんか普通に大して何も考えずに書いても同じくらいのスピードが出ました。(先程の問題の提出: std (18 ms), 自作 (18 ms) まあ速すぎて違いよくわかりませんはそうかもです。)

というわけで、みなさまも大して何も考えずに二分ヒープを実装しましょう。意外と適当にやっても別に遅くもなく、不便でもないです。

オススメのインターフェース

こういう感じで、const generics を用いて min-heap, max-heap が簡単に使えるようにするとよいです。

pub type MinHeap<T> = Heap<T, { HeapType::Min as usize }>;
pub type MaxHeap<T> = Heap<T, { HeapType::Max as usize }>;

pub enum HeapType {
    Min = 0,
    Max = 1,
}

また実際のアルゴリズム中の分岐は、次のように毎回 if で分岐するのがオススメです。なぜならこう書くと具体的に、min-heap の場合はこう、max-heap の場合はこう、というのがまさに並列されて書かれていてわかりやすいと私は考えているからです。

なおこのコードはコンパイル時に if 0 == 0 とかに変換されるので、最適化で消えるはずです。

impl<T: Ord, const TYPE: usize> Heap<T, TYPE> {
    pub fn push(&mut self, item: T) {
        let mut i = self.items.len();
        self.items.push(item);
        while i != 0 {
            let p = (i - 1) / 2;
            if {
                if TYPE == HeapType::Min as usize {
                    self.items[p] <= self.items[i]
                } else {
                    self.items[p] >= self.items[i]
                }
            } {
                break;
            }
            self.items.swap(p, i);
            i = p;
        }
    }
}

ちなみに二分ヒープの完全なコードは差し上げません。やーい釣られてやんの！

代案として、次のような感じでトレイトを定義するというのがありますが、あまりおすすめしません。要素の比較を比較をしたいときに毎回この形 (たとえば prioritize_over) への言い換えを行う必要が生じて、分かりづらいからです。コードを長くしてでも if TYPE == 方式のほうが良いと思います。そもそもトレイトは型定義パートにボイラープレートが増えるので総コード長はあまり変わらないというお話もあります。

pub trait HeapType {
    fn prioritize_over<T: Ord>(left: &T, right: &T) -> bool;
}

enum Min {}
impl HeapType for Min {
    fn prioritize_over<T: Ord>(left: &T, right: &T) -> bool {
        left < right
    }
}

enum Max {}
impl HeapType for Max {
    fn prioritize_over<T: Ord>(left: &T, right: &T) -> bool {
        left > right
    }
}

先程の問題のコード

自作すればこのように、対称性の高いコードになって読みやすいです！ これに Reversed や .0 がモリモリ入ると思うと悲しいですね。

use itertools::Itertools;
use proconio::input;

fn main() {
    input! {
        n: usize,
        lr: [(i64, i64); n],
    }
    let mut left = MaxHeap::new();
    let mut right = MinHeap::new();
    let mut left_offset = 0;
    let mut right_offset = 0;
    let mut min = 0;
    left.push(lr[0].0);
    right.push(lr[0].0);
    for ((l0, r0), (l1, r1)) in lr.iter().copied().tuple_windows() {
        left_offset -= r1 - l1;
        right_offset += r0 - l0;
        let x = l1;
        if x < left.peek().unwrap() + left_offset {
            min += left.peek().unwrap() + left_offset - x;
            left.push(x - left_offset);
            left.push(x - left_offset);
            right.push(left.pop().unwrap() + left_offset - right_offset);
        } else if right.peek().unwrap() + right_offset < x {
            min += x - (right.peek().unwrap() + right_offset);
            right.push(x - right_offset);
            right.push(x - right_offset);
            left.push(right.pop().unwrap() + right_offset - left_offset);
        } else {
            left.push(x - left_offset);
            right.push(x - right_offset);
        }
    }
    println!("{min}");
}

ちなみに Nightly なら？

adt-const-params を使うことでこれができてしまいます。楽しみですね。

pub type MinHeap<T> = Heap<T, HeapType::Min>;
pub type MaxHeap<T> = Heap<T, HeapType::Max>;