std.Treapをordered_mapとして振る舞わせるTips

はじめに

zig言語の標準ライブラリには、木構造を扱う型として、std.Treapが提供されている。
順序集合(ordered setとしての使用が主な目的となるが、ひと工夫することで順序マップ(ordered map)として扱う方法が見い出せたため、記録として残す。

方法

実装依存の方法のため、将来のバージョンで動かなくなるかもしれないことに注意。
確認は、2023-11-02時点のnightlyで行った。

std.Treap型は、以下のように定義されている。

pub fn Treap(comptime Key: type, comptime compareFn: anytype) { ... }

Keyは集合の要素となる値の型、compareFnはコンパレータで、以下のインターフェースを想定している。

fn compare(a: Key, b: Key) std.math.Order { ... }

コンパレータのインターフェースはキー同士を比較し、その順序を返すものとなっている。

Keyには任意の型を指定できるため、以下のような構造体とコンパレータを指定するで、真のキーと値とのマップを模倣できる。

const Data = struct { 
    key: SomeKeyType, // 真のキー
    value: SomeValueType, // キーにマップしたい値
};

const Comparator = struct {
    pub fn compare (a: Data, b: Data) std.math.Order {
        // a.keyとb.keyの比較結果を返す。
    }

エントリを追加する際は、以下のように行う。

• Treap.getEntryForに対して、真のキーのみを詰めた値で探索させる。

var tree = Treap(....){};

// 探索の際、値は不要のためundefinedを詰めておく
var entry = tree.getEntryFor(Data{ .key = ..., .value = undefined, });

• 探索結果の、TreapのEntryのキーを本来セットしたいものに差し替え、ツリーを更新する。

if (entry.node == null) {
  // 追加したい真のキーと値を割り当てる
  // entry.setメソッドでは、entry.keyを使用して木構造を再構成するため
  entry.key = Data{ .key = ..., .value = ... };

  // Nodeを作成する
  var node = try allocator.create(Treap(...).Node);
  node.* = undefined;

  // ツリーの要素を更新する
  entry.set(node);
}

zig言語バージョン、0.12系で、Treapに対するイテレータを取得するinorderIteratorメソッドが追加されたため、このイテレータを使用することで真のキーでソートされた結果を取り出すことができる。

var iter = tree.inorderIterator();
while (iter.next()) |node| {
  // node.key.value はnode.key.keyでソートされた順に取り出せる
}