std.TreapのIteratorを作る

はじめに

Zig言語の標準ライブラリには、二分木を扱うための型としてTreapが提供されている。

しかしこのTreap型には残念ながらIteratorが提供されていないため、小さい順に取り出すのがちと面倒い。

必要に駆られたので頑張って作ってみた。

実装

以下の解説記事をにらめっこしながら実装した。

以下のコードが出来上がったもの

fn TreapIterator(comptime TTreap: type, comptime TKey: type) type {
    return struct {
        node: ?*TTreap.Node,
        
        // 初期化時に最小の要素を指しておく
        // 最小の要素はTreap.getMin()で取得できる
        fn init(tree: TTreap) @This() {
            return .{
                .node = tree.getMin(),
            };
        }
        
        // 小さい順にツリーをたどって要素を返していく
        pub fn next(self: *@This()) ?TKey {
            return if (self.node) |current| find_next: {
                var n: ?*TTreap.Node = null;
                defer self.node = n;

                if (current.children[1]) |n0| {
                    n = n0;
                    while (n) |n1| {
                        n = n1.children[0] orelse break;
                    }
                }
                else {
                    n = current.parent;
                    var n0 = self.node;
                    while (n) |n1| {
                        if (n1.children[1] != n0) break;

                        n = n1.parent;
                        n0 = n1;
                    }
                }

                break :find_next current.*.key;
            }
            else null;
        }
    };
}

使用例

pub fn main() void {
    const IntTreap = std.Treap(i64, std.math.order);
    const Node = IntTreap.Node;
    
    var t = IntTreap{};
    var nodes: [10]Node = undefined;

    var rand_seed = std.rand.DefaultPrng.init(@intCast(u64, std.time.milliTimestamp()));
    var rand_gen = rand_seed.random();
    
    std.debug.print("------------------------------\n", .{});
    std.debug.print("> Generate Sequences\n", .{});
    std.debug.print("------------------------------\n", .{});

    var i: usize = 0;
    while (i < 10): (i += 1) {
        var x = rand_gen.int(i64);
        var entry = t.getEntryFor(x);

        nodes[i].key = x;
        entry.set(&nodes[i]);

        std.debug.print("{}\n", .{ nodes[i].key });
    }

    std.debug.print("------------------------------\n", .{});
    std.debug.print("> Iterate Sequences\n", .{});
    std.debug.print("------------------------------\n", .{});

    var it = TreapIterator(IntTreap, i64).init(t);

    while (it.next()) |x| {
        std.debug.print("{}\n", .{ x });
    }
}

実行結果（例）

% zig build run
------------------------------
> Generate Sequences
------------------------------
7506171842220577781
5512948215387453786
-1282317308755977892
5040489507533034472
2766790032025837170
1634193634318532548
-3883475461565456234
-2156708153179984805
524639672238038339
-2475206045401187818
------------------------------
> Iterate Sequences
------------------------------
-3883475461565456234
-2475206045401187818
-2156708153179984805
-1282317308755977892
524639672238038339
1634193634318532548
2766790032025837170
5040489507533034472
5512948215387453786
7506171842220577781

2023-12-06 追記

zig言語のバージョン0.12系で、晴れてTreapのイテレータが標準搭載された。