authorgravatar for john.schmidt.h@gmail.comJohn Schmidt <john.schmidt.h@gmail.com> 2019-02-26 20:44:27+01:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2019-03-08 14:23:05-05:00
logdd34c217791a0949ef2b0531ce8e8d4691239c45
treee3863a75eae58531eae74ab0996490ff36515a3e
parentcddc77dd97fa222fc72e4b3955eabb18c36c8617

PriorityQueue: add bulk insertion methods

1 files changed, 59 insertions(+), 21 deletions(-)

std/priority_queue.zig+59-21
......@@ -28,7 +28,10 @@ pub fn PriorityQueue(comptime T: type) type {
2828
2929 pub fn add(self: *Self, elem: T) !void {
3030 try ensureCapacity(self, self.len + 1);
31 addUnchecked(self, elem);
32 }
3133
34 fn addUnchecked(self: *Self, elem: T) void {
3235 self.items[self.len] = elem;
3336 var child_index = self.len;
3437 while (child_index > 0) {
......@@ -45,6 +48,13 @@ pub fn PriorityQueue(comptime T: type) type {
4548 self.len += 1;
4649 }
4750
51 pub fn addSlice(self: *Self, items: []const T) !void {
52 try self.ensureCapacity(self.len + items.len);
53 for (items) |e| {
54 self.addUnchecked(e);
55 }
56 }
57
4858 pub fn peek(self: *Self) ?T {
4959 return if (self.len > 0) self.items[0] else null;
5060 }
......@@ -58,7 +68,7 @@ pub fn PriorityQueue(comptime T: type) type {
5868 const last = self.items[self.len - 1];
5969 self.items[0] = last;
6070 self.len -= 1;
61 siftDown(self);
71 siftDown(self, 0);
6272 return first;
6373 }
6474
......@@ -70,8 +80,8 @@ pub fn PriorityQueue(comptime T: type) type {
7080 return self.items.len;
7181 }
7282
73 fn siftDown(self: *Self) void {
74 var index: usize = 0;
83 fn siftDown(self: *Self, start_index: usize) void {
84 var index = start_index;
7585 const half = self.len >> 1;
7686 while (true) {
7787 var left_index = (index << 1) + 1;
......@@ -106,6 +116,24 @@ pub fn PriorityQueue(comptime T: type) type {
106116 }
107117 }
108118
119 /// PriorityQueue takes ownership of the passed in slice. The slice must have been
120 /// allocated with `allocator`.
121 /// Deinitialize with `deinit`.
122 pub fn fromOwnedSlice(allocator: *Allocator, compareFn: fn (a: T, b: T) bool, items: []T) Self {
123 var queue = Self{
124 .items = items,
125 .len = items.len,
126 .allocator = allocator,
127 .compareFn = compareFn,
128 };
129 const half = (queue.len >> 1) - 1;
130 var i: usize = 0;
131 while (i <= half) : (i += 1) {
132 queue.siftDown(half - i);
133 }
134 return queue;
135 }
136
109137 pub fn ensureCapacity(self: *Self, new_capacity: usize) !void {
110138 var better_capacity = self.capacity();
111139 if (better_capacity >= new_capacity) return;
......@@ -252,24 +280,34 @@ test "std.PriorityQueue: sift up with odd indices" {
252280 try queue.add(e);
253281 }
254282
255 expectEqual(u32(1), queue.remove());
256 expectEqual(u32(2), queue.remove());
257 expectEqual(u32(5), queue.remove());
258 expectEqual(u32(6), queue.remove());
259 expectEqual(u32(7), queue.remove());
260 expectEqual(u32(7), queue.remove());
261 expectEqual(u32(11), queue.remove());
262 expectEqual(u32(12), queue.remove());
263 expectEqual(u32(13), queue.remove());
264 expectEqual(u32(14), queue.remove());
265 expectEqual(u32(15), queue.remove());
266 expectEqual(u32(15), queue.remove());
267 expectEqual(u32(16), queue.remove());
268 expectEqual(u32(21), queue.remove());
269 expectEqual(u32(22), queue.remove());
270 expectEqual(u32(24), queue.remove());
271 expectEqual(u32(24), queue.remove());
272 expectEqual(u32(25), queue.remove());
283 const sorted_items = []u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
284 for (sorted_items) |e| {
285 expectEqual(e, queue.remove());
286 }
287}
288
289test "std.PriorityQueue: addSlice" {
290 var queue = PQ.init(debug.global_allocator, lessThan);
291 defer queue.deinit();
292 const items = []u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
293 try queue.addSlice(items[0..]);
294
295 const sorted_items = []u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
296 for (sorted_items) |e| {
297 expectEqual(e, queue.remove());
298 }
299}
300
301test "std.PriorityQueue: fromOwnedSlice" {
302 const items = []u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
303 const heap_items = try std.mem.dupe(debug.global_allocator, u32, items[0..]);
304 var queue = PQ.fromOwnedSlice(debug.global_allocator, lessThan, heap_items[0..]);
305 defer queue.deinit();
306
307 const sorted_items = []u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
308 for (sorted_items) |e| {
309 expectEqual(e, queue.remove());
310 }
273311}
274312
275313test "std.PriorityQueue: add and remove max heap" {