| ... | ... | @@ -190,6 +190,31 @@ pub fn PriorityQueue(comptime T: type) type { |
| 190 | 190 | self.len = new_len; |
| 191 | 191 | } |
| 192 | 192 | |
| 193 | fn binarySearch(items: []const T, target: T) usize { |
| 194 | var left: usize = 0; |
| 195 | var right: usize = items.len-1; |
| 196 | |
| 197 | while(left <= right) { |
| 198 | const mid = left + (right - left) / 2; |
| 199 | if (items[mid] == target) { |
| 200 | return mid; |
| 201 | } else if (items[mid] < target) { |
| 202 | left= mid+1; |
| 203 | } else { |
| 204 | right= mid-1; |
| 205 | } |
| 206 | } |
| 207 | |
| 208 | return 0; |
| 209 | } |
| 210 | |
| 211 | pub fn update(self: *Self, elem: T, new_elem: T) !void { |
| 212 | var update_index: usize = binarySearch(self.items, elem); |
| 213 | assert (update_index >= 0 and update_index < self.items.len); |
| 214 | _ = self.removeIndex(update_index); |
| 215 | try self.add(new_elem); |
| 216 | } |
| 217 | |
| 193 | 218 | pub const Iterator = struct { |
| 194 | 219 | queue: *PriorityQueue(T), |
| 195 | 220 | count: usize, |
| ... | ... | @@ -437,3 +462,66 @@ test "std.PriorityQueue: iterator while empty" { |
| 437 | 462 | |
| 438 | 463 | expectEqual(it.next(), null); |
| 439 | 464 | } |
| 465 | |
| 466 | test "std.PriorityQueue: update min heap" { |
| 467 | var queue = PQ.init(testing.allocator, lessThan); |
| 468 | defer queue.deinit(); |
| 469 | |
| 470 | try queue.add(55); |
| 471 | try queue.add(44); |
| 472 | try queue.add(11); |
| 473 | try queue.update(55, 5); |
| 474 | try queue.update(44, 4); |
| 475 | try queue.update(11, 1); |
| 476 | expectEqual(@as(u32, 1), queue.remove()); |
| 477 | expectEqual(@as(u32, 4), queue.remove()); |
| 478 | expectEqual(@as(u32, 5), queue.remove()); |
| 479 | } |
| 480 | |
| 481 | |
| 482 | test "std.PriorityQueue: update same min heap" { |
| 483 | var queue = PQ.init(testing.allocator, lessThan); |
| 484 | defer queue.deinit(); |
| 485 | |
| 486 | try queue.add(1); |
| 487 | try queue.add(1); |
| 488 | try queue.add(2); |
| 489 | try queue.add(2); |
| 490 | try queue.update(1, 5); |
| 491 | try queue.update(2, 4); |
| 492 | expectEqual(@as(u32, 1), queue.remove()); |
| 493 | expectEqual(@as(u32, 2), queue.remove()); |
| 494 | expectEqual(@as(u32, 4), queue.remove()); |
| 495 | expectEqual(@as(u32, 5), queue.remove()); |
| 496 | } |
| 497 | |
| 498 | test "std.PriorityQueue: update max heap" { |
| 499 | var queue = PQ.init(testing.allocator, greaterThan); |
| 500 | defer queue.deinit(); |
| 501 | |
| 502 | try queue.add(55); |
| 503 | try queue.add(44); |
| 504 | try queue.add(11); |
| 505 | try queue.update(55, 5); |
| 506 | try queue.update(44, 1); |
| 507 | try queue.update(11, 4); |
| 508 | expectEqual(@as(u32, 5), queue.remove()); |
| 509 | expectEqual(@as(u32, 4), queue.remove()); |
| 510 | expectEqual(@as(u32, 1), queue.remove()); |
| 511 | } |
| 512 | |
| 513 | test "std.PriorityQueue: update same max heap" { |
| 514 | var queue = PQ.init(testing.allocator, greaterThan); |
| 515 | defer queue.deinit(); |
| 516 | |
| 517 | try queue.add(1); |
| 518 | try queue.add(1); |
| 519 | try queue.add(2); |
| 520 | try queue.add(2); |
| 521 | try queue.update(1, 5); |
| 522 | try queue.update(2, 4); |
| 523 | expectEqual(@as(u32, 5), queue.remove()); |
| 524 | expectEqual(@as(u32, 4), queue.remove()); |
| 525 | expectEqual(@as(u32, 2), queue.remove()); |
| 526 | expectEqual(@as(u32, 1), queue.remove()); |
| 527 | } |
| | \ No newline at end of file |