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 {...@@ -28,7 +28,10 @@ pub fn PriorityQueue(comptime T: type) type {
2828
29 pub fn add(self: *Self, elem: T) !void {29 pub fn add(self: *Self, elem: T) !void {
30 try ensureCapacity(self, self.len + 1);30 try ensureCapacity(self, self.len + 1);
31 addUnchecked(self, elem);
32 }
3133
34 fn addUnchecked(self: *Self, elem: T) void {
32 self.items[self.len] = elem;35 self.items[self.len] = elem;
33 var child_index = self.len;36 var child_index = self.len;
34 while (child_index > 0) {37 while (child_index > 0) {
...@@ -45,6 +48,13 @@ pub fn PriorityQueue(comptime T: type) type {...@@ -45,6 +48,13 @@ pub fn PriorityQueue(comptime T: type) type {
45 self.len += 1;48 self.len += 1;
46 }49 }
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
48 pub fn peek(self: *Self) ?T {58 pub fn peek(self: *Self) ?T {
49 return if (self.len > 0) self.items[0] else null;59 return if (self.len > 0) self.items[0] else null;
50 }60 }
...@@ -58,7 +68,7 @@ pub fn PriorityQueue(comptime T: type) type {...@@ -58,7 +68,7 @@ pub fn PriorityQueue(comptime T: type) type {
58 const last = self.items[self.len - 1];68 const last = self.items[self.len - 1];
59 self.items[0] = last;69 self.items[0] = last;
60 self.len -= 1;70 self.len -= 1;
61 siftDown(self);71 siftDown(self, 0);
62 return first;72 return first;
63 }73 }
6474
...@@ -70,8 +80,8 @@ pub fn PriorityQueue(comptime T: type) type {...@@ -70,8 +80,8 @@ pub fn PriorityQueue(comptime T: type) type {
70 return self.items.len;80 return self.items.len;
71 }81 }
7282
73 fn siftDown(self: *Self) void {83 fn siftDown(self: *Self, start_index: usize) void {
74 var index: usize = 0;84 var index = start_index;
75 const half = self.len >> 1;85 const half = self.len >> 1;
76 while (true) {86 while (true) {
77 var left_index = (index << 1) + 1;87 var left_index = (index << 1) + 1;
...@@ -106,6 +116,24 @@ pub fn PriorityQueue(comptime T: type) type {...@@ -106,6 +116,24 @@ pub fn PriorityQueue(comptime T: type) type {
106 }116 }
107 }117 }
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
109 pub fn ensureCapacity(self: *Self, new_capacity: usize) !void {137 pub fn ensureCapacity(self: *Self, new_capacity: usize) !void {
110 var better_capacity = self.capacity();138 var better_capacity = self.capacity();
111 if (better_capacity >= new_capacity) return;139 if (better_capacity >= new_capacity) return;
...@@ -252,24 +280,34 @@ test "std.PriorityQueue: sift up with odd indices" {...@@ -252,24 +280,34 @@ test "std.PriorityQueue: sift up with odd indices" {
252 try queue.add(e);280 try queue.add(e);
253 }281 }
254282
255 expectEqual(u32(1), 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 };
256 expectEqual(u32(2), queue.remove());284 for (sorted_items) |e| {
257 expectEqual(u32(5), queue.remove());285 expectEqual(e, queue.remove());
258 expectEqual(u32(6), queue.remove());286 }
259 expectEqual(u32(7), queue.remove());287}
260 expectEqual(u32(7), queue.remove());288
261 expectEqual(u32(11), queue.remove());289test "std.PriorityQueue: addSlice" {
262 expectEqual(u32(12), queue.remove());290 var queue = PQ.init(debug.global_allocator, lessThan);
263 expectEqual(u32(13), queue.remove());291 defer queue.deinit();
264 expectEqual(u32(14), queue.remove());292 const items = []u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
265 expectEqual(u32(15), queue.remove());293 try queue.addSlice(items[0..]);
266 expectEqual(u32(15), queue.remove());294
267 expectEqual(u32(16), queue.remove());295 const sorted_items = []u32{ 1, 2, 5, 6, 7, 7, 11, 12, 13, 14, 15, 15, 16, 21, 22, 24, 24, 25 };
268 expectEqual(u32(21), queue.remove());296 for (sorted_items) |e| {
269 expectEqual(u32(22), queue.remove());297 expectEqual(e, queue.remove());
270 expectEqual(u32(24), queue.remove());298 }
271 expectEqual(u32(24), queue.remove());299}
272 expectEqual(u32(25), queue.remove());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 }
273}311}
274312
275test "std.PriorityQueue: add and remove max heap" {313test "std.PriorityQueue: add and remove max heap" {