authorgravatar for Validark@pm.meNiles Salter <Validark@pm.me> 2023-06-23 13:18:56-06:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2023-07-24 12:58:19-07:00
log8b1976cab55d04a138ccb102ca2008558e34bd7d
tree83c7cbdd826d63cd9b2c7dfc85e0638040282f40
parentf1bd59876872aad2be61a322eb07812b267a9997

[priority_queue] Simplify sifting & fix edge case

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

lib/std/priority_queue.zig+21-41
......@@ -51,18 +51,16 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
5151 }
5252
5353 fn siftUp(self: *Self, start_index: usize) void {
54 const child = self.items[start_index];
5455 var child_index = start_index;
5556 while (child_index > 0) {
56 var parent_index = ((child_index - 1) >> 1);
57 const child = self.items[child_index];
57 const parent_index = ((child_index - 1) >> 1);
5858 const parent = self.items[parent_index];
59
6059 if (compareFn(self.context, child, parent) != .lt) break;
61
62 self.items[parent_index] = child;
6360 self.items[child_index] = parent;
6461 child_index = parent_index;
6562 }
63 self.items[child_index] = child;
6664 }
6765
6866 /// Add each element in `items` to the queue.
......@@ -128,61 +126,43 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
128126 return self.items.len;
129127 }
130128
131 fn siftDown(self: *Self, start_index: usize) void {
132 var index = start_index;
133 const half = self.len >> 1;
129 fn siftDown(self: *Self, target_index: usize) void {
130 const target_element = self.items[target_index];
131 var index = target_index;
134132 while (true) {
135 var left_index = (index << 1) + 1;
136 var right_index = left_index + 1;
137 var left = if (left_index < self.len) self.items[left_index] else null;
138 var right = if (right_index < self.len) self.items[right_index] else null;
139
140 var smallest_index = index;
141 var smallest = self.items[index];
142
143 if (left) |e| {
144 if (compareFn(self.context, e, smallest) == .lt) {
145 smallest_index = left_index;
146 smallest = e;
147 }
148 }
133 var lesser_child_i = (std.math.mul(usize, index, 2) catch break) | 1;
134 if (!(lesser_child_i < self.len)) break;
149135
150 if (right) |e| {
151 if (compareFn(self.context, e, smallest) == .lt) {
152 smallest_index = right_index;
153 smallest = e;
154 }
136 const next_child_i = lesser_child_i + 1;
137 if (next_child_i < self.len and compareFn(self.context, self.items[next_child_i], self.items[lesser_child_i]) == .lt) {
138 lesser_child_i = next_child_i;
155139 }
156140
157 if (smallest_index == index) return;
141 if (compareFn(self.context, target_element, self.items[lesser_child_i]) == .lt) break;
158142
159 self.items[smallest_index] = self.items[index];
160 self.items[index] = smallest;
161 index = smallest_index;
162
163 if (index >= half) return;
143 self.items[index] = self.items[lesser_child_i];
144 index = lesser_child_i;
164145 }
146 self.items[index] = target_element;
165147 }
166148
167149 /// PriorityQueue takes ownership of the passed in slice. The slice must have been
168150 /// allocated with `allocator`.
169151 /// Deinitialize with `deinit`.
170152 pub fn fromOwnedSlice(allocator: Allocator, items: []T, context: Context) Self {
171 var queue = Self{
153 var self = Self{
172154 .items = items,
173155 .len = items.len,
174156 .allocator = allocator,
175157 .context = context,
176158 };
177159
178 if (queue.len <= 1) return queue;
179
180 const half = (queue.len >> 1) - 1;
181 var i: usize = 0;
182 while (i <= half) : (i += 1) {
183 queue.siftDown(half - i);
160 var i = self.len >> 1;
161 while (i > 0) {
162 i -= 1;
163 self.siftDown(i);
184164 }
185 return queue;
165 return self;
186166 }
187167
188168 /// Ensure that the queue can fit at least `new_capacity` items.