authorgravatar for john.schmidt.h@gmail.comJohn Schmidt <john.schmidt.h@gmail.com> 2019-02-25 20:25:21+01:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2019-02-26 13:21:11-05:00
log0eed72d6874b1705405df0d7975c66a2178c7106
treea286d7f66ab458aba9708af93532e341b22c1c58
parent6fd28164316185676ca3337f07c510b409cb6687

Add priority queue

3 files changed, 335 insertions(+), 1 deletions(-)

CMakeLists.txt+1
......@@ -609,6 +609,7 @@ set(ZIG_STD_FILES
609609 "os/windows/util.zig"
610610 "os/zen.zig"
611611 "pdb.zig"
612 "priority_queue.zig"
612613 "rand/index.zig"
613614 "rand/ziggurat.zig"
614615 "segmented_list.zig"
std/index.zig+3-1
......@@ -9,6 +9,7 @@ pub const DynLib = @import("dynamic_library.zig").DynLib;
99pub const HashMap = @import("hash_map.zig").HashMap;
1010pub const LinkedList = @import("linked_list.zig").LinkedList;
1111pub const Mutex = @import("mutex.zig").Mutex;
12pub const PriorityQueue = @import("priority_queue.zig").PriorityQueue;
1213pub const StaticallyInitializedMutex = @import("statically_initialized_mutex.zig").StaticallyInitializedMutex;
1314pub const SegmentedList = @import("segmented_list.zig").SegmentedList;
1415pub const SpinLock = @import("spinlock.zig").SpinLock;
......@@ -59,7 +60,7 @@ test "std" {
5960 _ = @import("statically_initialized_mutex.zig");
6061 _ = @import("segmented_list.zig");
6162 _ = @import("spinlock.zig");
62
63
6364 _ = @import("base64.zig");
6465 _ = @import("build.zig");
6566 _ = @import("c/index.zig");
......@@ -85,6 +86,7 @@ test "std" {
8586 _ = @import("net.zig");
8687 _ = @import("os/index.zig");
8788 _ = @import("pdb.zig");
89 _ = @import("priority_queue.zig");
8890 _ = @import("rand/index.zig");
8991 _ = @import("sort.zig");
9092 _ = @import("testing.zig");
std/priority_queue.zig created+331
......@@ -0,0 +1,331 @@
1const std = @import("index.zig");
2const Allocator = std.mem.Allocator;
3const debug = std.debug;
4const expect = std.testing.expect;
5const expectEqual = std.testing.expectEqual;
6
7pub fn PriorityQueue(comptime T: type) type {
8 return struct {
9 const Self = @This();
10
11 items: []T,
12 len: usize,
13 allocator: *Allocator,
14 compareFn: fn (a: T, b: T) bool,
15
16 pub fn init(allocator: *Allocator, compareFn: fn (a: T, b: T) bool) Self {
17 return Self{
18 .items = []T{},
19 .len = 0,
20 .allocator = allocator,
21 .compareFn = compareFn,
22 };
23 }
24
25 pub fn deinit(self: Self) void {
26 self.allocator.free(self.items);
27 }
28
29 pub fn add(self: *Self, elem: T) !void {
30 try ensureCapacity(self, self.len + 1);
31
32 self.items[self.len] = elem;
33 var child_index = self.len;
34 while (child_index > 0) {
35 var parent_index = ((child_index - 1) >> 1);
36 const child = self.items[child_index];
37 const parent = self.items[parent_index];
38
39 if (!self.compareFn(child, parent)) break;
40
41 self.items[parent_index] = child;
42 self.items[child_index] = parent;
43 child_index = parent_index;
44 }
45 self.len += 1;
46 }
47
48 pub fn peek(self: *Self) ?T {
49 return if (self.len > 0) self.items[0] else null;
50 }
51
52 pub fn removeOrNull(self: *Self) ?T {
53 return if (self.len > 0) self.remove() else null;
54 }
55
56 pub fn remove(self: *Self) T {
57 const first = self.items[0];
58 const last = self.items[self.len - 1];
59 self.items[0] = last;
60 self.len -= 1;
61 siftDown(self);
62 return first;
63 }
64
65 pub fn count(self: Self) usize {
66 return self.len;
67 }
68
69 pub fn capacity(self: Self) usize {
70 return self.items.len;
71 }
72
73 fn siftDown(self: *Self) void {
74 var index: usize = 0;
75 const half = self.len >> 1;
76 while (true) {
77 var left_index = (index << 1) + 1;
78 var right_index = left_index + 1;
79 var left = if (left_index < self.len) self.items[left_index] else null;
80 var right = if (right_index < self.len) self.items[right_index] else null;
81
82 var smallest_index = index;
83 var smallest = self.items[index];
84
85 if (left) |e| {
86 if (self.compareFn(e, smallest)) {
87 smallest_index = left_index;
88 smallest = e;
89 }
90 }
91
92 if (right) |e| {
93 if (self.compareFn(e, smallest)) {
94 smallest_index = right_index;
95 smallest = e;
96 }
97 }
98
99 if (smallest_index == index) return;
100
101 self.items[smallest_index] = self.items[index];
102 self.items[index] = smallest;
103 index = smallest_index;
104
105 if (index >= half) return;
106 }
107 }
108
109 pub fn ensureCapacity(self: *Self, new_capacity: usize) !void {
110 var better_capacity = self.capacity();
111 if (better_capacity >= new_capacity) return;
112 while (true) {
113 better_capacity += better_capacity / 2 + 8;
114 if (better_capacity >= new_capacity) break;
115 }
116 self.items = try self.allocator.realloc(T, self.items, better_capacity);
117 }
118
119 pub fn resize(self: *Self, new_len: usize) !void {
120 try self.ensureCapacity(new_len);
121 self.len = new_len;
122 }
123
124 pub fn shrink(self: *Self, new_len: usize) void {
125 assert(new_len <= self.len);
126 self.len = new_len;
127 }
128
129 const Iterator = struct {
130 queue: *PriorityQueue(T),
131 count: usize,
132
133 fn next(it: *Iterator) ?T {
134 if (it.count > it.queue.len - 1) return null;
135 const out = it.count;
136 it.count += 1;
137 return it.queue.items[out];
138 }
139
140 fn reset(it: *Iterator) void {
141 it.count = 0;
142 }
143 };
144
145 pub fn iterator(self: *Self) Iterator {
146 return Iterator{
147 .queue = self,
148 .count = 0,
149 };
150 }
151
152 fn dump(self: *Self) void {
153 warn("{{ ");
154 warn("items: ");
155 for (self.items) |e, i| {
156 if (i >= self.len) break;
157 warn("{}, ", e);
158 }
159 warn("array: ");
160 for (self.items) |e, i| {
161 warn("{}, ", e);
162 }
163 warn("len: {} ", self.len);
164 warn("capacity: {}", self.capacity());
165 warn(" }}\n");
166 }
167 };
168}
169
170fn lessThan(a: u32, b: u32) bool {
171 return a < b;
172}
173
174fn greaterThan(a: u32, b: u32) bool {
175 return a > b;
176}
177
178const PQ = PriorityQueue(u32);
179
180test "std.PriorityQueue: add and remove min heap" {
181 var queue = PQ.init(debug.global_allocator, lessThan);
182 defer queue.deinit();
183
184 try queue.add(54);
185 try queue.add(12);
186 try queue.add(7);
187 try queue.add(23);
188 try queue.add(25);
189 try queue.add(13);
190 expectEqual(u32(7), queue.remove());
191 expectEqual(u32(12), queue.remove());
192 expectEqual(u32(13), queue.remove());
193 expectEqual(u32(23), queue.remove());
194 expectEqual(u32(25), queue.remove());
195 expectEqual(u32(54), queue.remove());
196}
197
198test "std.PriorityQueue: add and remove same min heap" {
199 var queue = PQ.init(debug.global_allocator, lessThan);
200 defer queue.deinit();
201
202 try queue.add(1);
203 try queue.add(1);
204 try queue.add(2);
205 try queue.add(2);
206 try queue.add(1);
207 try queue.add(1);
208 expectEqual(u32(1), queue.remove());
209 expectEqual(u32(1), queue.remove());
210 expectEqual(u32(1), queue.remove());
211 expectEqual(u32(1), queue.remove());
212 expectEqual(u32(2), queue.remove());
213 expectEqual(u32(2), queue.remove());
214}
215
216test "std.PriorityQueue: removeOrNull on empty" {
217 var queue = PQ.init(debug.global_allocator, lessThan);
218 defer queue.deinit();
219
220 expect(queue.removeOrNull() == null);
221}
222
223test "std.PriorityQueue: edge case 3 elements" {
224 var queue = PQ.init(debug.global_allocator, lessThan);
225 defer queue.deinit();
226
227 try queue.add(9);
228 try queue.add(3);
229 try queue.add(2);
230 expectEqual(u32(2), queue.remove());
231 expectEqual(u32(3), queue.remove());
232 expectEqual(u32(9), queue.remove());
233}
234
235test "std.PriorityQueue: peek" {
236 var queue = PQ.init(debug.global_allocator, lessThan);
237 defer queue.deinit();
238
239 expect(queue.peek() == null);
240 try queue.add(9);
241 try queue.add(3);
242 try queue.add(2);
243 expectEqual(u32(2), queue.peek().?);
244 expectEqual(u32(2), queue.peek().?);
245}
246
247test "std.PriorityQueue: sift up with odd indices" {
248 var queue = PQ.init(debug.global_allocator, lessThan);
249 defer queue.deinit();
250 const items = []u32{ 15, 7, 21, 14, 13, 22, 12, 6, 7, 25, 5, 24, 11, 16, 15, 24, 2, 1 };
251 for (items) |e| {
252 try queue.add(e);
253 }
254
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());
273}
274
275test "std.PriorityQueue: add and remove max heap" {
276 var queue = PQ.init(debug.global_allocator, greaterThan);
277 defer queue.deinit();
278
279 try queue.add(54);
280 try queue.add(12);
281 try queue.add(7);
282 try queue.add(23);
283 try queue.add(25);
284 try queue.add(13);
285 expectEqual(u32(54), queue.remove());
286 expectEqual(u32(25), queue.remove());
287 expectEqual(u32(23), queue.remove());
288 expectEqual(u32(13), queue.remove());
289 expectEqual(u32(12), queue.remove());
290 expectEqual(u32(7), queue.remove());
291}
292
293test "std.PriorityQueue: add and remove same max heap" {
294 var queue = PQ.init(debug.global_allocator, greaterThan);
295 defer queue.deinit();
296
297 try queue.add(1);
298 try queue.add(1);
299 try queue.add(2);
300 try queue.add(2);
301 try queue.add(1);
302 try queue.add(1);
303 expectEqual(u32(2), queue.remove());
304 expectEqual(u32(2), queue.remove());
305 expectEqual(u32(1), queue.remove());
306 expectEqual(u32(1), queue.remove());
307 expectEqual(u32(1), queue.remove());
308 expectEqual(u32(1), queue.remove());
309}
310
311test "std.PriorityQueue: iterator" {
312 var queue = PQ.init(debug.global_allocator, lessThan);
313 var map = std.AutoHashMap(u32, void).init(debug.global_allocator);
314 defer {
315 queue.deinit();
316 map.deinit();
317 }
318
319 const items = []u32{ 54, 12, 7, 23, 25, 13 };
320 for (items) |e| {
321 _ = try queue.add(e);
322 _ = try map.put(e, {});
323 }
324
325 var it = queue.iterator();
326 while (it.next()) |e| {
327 _ = map.remove(e);
328 }
329
330 expectEqual(usize(0), map.count());
331}