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...@@ -51,18 +51,16 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
51 }51 }
5252
53 fn siftUp(self: *Self, start_index: usize) void {53 fn siftUp(self: *Self, start_index: usize) void {
54 const child = self.items[start_index];
54 var child_index = start_index;55 var child_index = start_index;
55 while (child_index > 0) {56 while (child_index > 0) {
56 var parent_index = ((child_index - 1) >> 1);57 const parent_index = ((child_index - 1) >> 1);
57 const child = self.items[child_index];
58 const parent = self.items[parent_index];58 const parent = self.items[parent_index];
59
60 if (compareFn(self.context, child, parent) != .lt) break;59 if (compareFn(self.context, child, parent) != .lt) break;
61
62 self.items[parent_index] = child;
63 self.items[child_index] = parent;60 self.items[child_index] = parent;
64 child_index = parent_index;61 child_index = parent_index;
65 }62 }
63 self.items[child_index] = child;
66 }64 }
6765
68 /// Add each element in `items` to the queue.66 /// Add each element in `items` to the queue.
...@@ -128,61 +126,43 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF...@@ -128,61 +126,43 @@ pub fn PriorityQueue(comptime T: type, comptime Context: type, comptime compareF
128 return self.items.len;126 return self.items.len;
129 }127 }
130128
131 fn siftDown(self: *Self, start_index: usize) void {129 fn siftDown(self: *Self, target_index: usize) void {
132 var index = start_index;130 const target_element = self.items[target_index];
133 const half = self.len >> 1;131 var index = target_index;
134 while (true) {132 while (true) {
135 var left_index = (index << 1) + 1;133 var lesser_child_i = (std.math.mul(usize, index, 2) catch break) | 1;
136 var right_index = left_index + 1;134 if (!(lesser_child_i < self.len)) break;
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 }
149135
150 if (right) |e| {136 const next_child_i = lesser_child_i + 1;
151 if (compareFn(self.context, e, smallest) == .lt) {137 if (next_child_i < self.len and compareFn(self.context, self.items[next_child_i], self.items[lesser_child_i]) == .lt) {
152 smallest_index = right_index;138 lesser_child_i = next_child_i;
153 smallest = e;
154 }
155 }139 }
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];143 self.items[index] = self.items[lesser_child_i];
160 self.items[index] = smallest;144 index = lesser_child_i;
161 index = smallest_index;
162
163 if (index >= half) return;
164 }145 }
146 self.items[index] = target_element;
165 }147 }
166148
167 /// PriorityQueue takes ownership of the passed in slice. The slice must have been149 /// PriorityQueue takes ownership of the passed in slice. The slice must have been
168 /// allocated with `allocator`.150 /// allocated with `allocator`.
169 /// Deinitialize with `deinit`.151 /// Deinitialize with `deinit`.
170 pub fn fromOwnedSlice(allocator: Allocator, items: []T, context: Context) Self {152 pub fn fromOwnedSlice(allocator: Allocator, items: []T, context: Context) Self {
171 var queue = Self{153 var self = Self{
172 .items = items,154 .items = items,
173 .len = items.len,155 .len = items.len,
174 .allocator = allocator,156 .allocator = allocator,
175 .context = context,157 .context = context,
176 };158 };
177159
178 if (queue.len <= 1) return queue;160 var i = self.len >> 1;
179161 while (i > 0) {
180 const half = (queue.len >> 1) - 1;162 i -= 1;
181 var i: usize = 0;163 self.siftDown(i);
182 while (i <= half) : (i += 1) {
183 queue.siftDown(half - i);
184 }164 }
185 return queue;165 return self;
186 }166 }
187167
188 /// Ensure that the queue can fit at least `new_capacity` items.168 /// Ensure that the queue can fit at least `new_capacity` items.