Keep only the top k objects according to their values.
| 457 | |
| 458 | |
| 459 | class BeamQueue: |
| 460 | """Keep only the top k objects according to their values.""" |
| 461 | |
| 462 | def __init__(self, max_size: int = 512): |
| 463 | self.queue = [] |
| 464 | self.max_size = max_size |
| 465 | |
| 466 | def add(self, node: object, val: float) -> None: |
| 467 | """Add a new node to this queue.""" |
| 468 | |
| 469 | if len(self.queue) < self.max_size: |
| 470 | self.queue.append((val, node)) |
| 471 | return |
| 472 | |
| 473 | # Find the minimum node: |
| 474 | min_idx, (min_val, _) = min(enumerate(self.queue), key=lambda x: x[1]) |
| 475 | |
| 476 | # replace it if the new node has higher value. |
| 477 | if val > min_val: |
| 478 | self.queue[min_idx] = (val, node) |
| 479 | |
| 480 | def __iter__(self): |
| 481 | for val, node in self.queue: |
| 482 | yield val, node |
| 483 | |
| 484 | def __len__(self) -> int: |
| 485 | return len(self.queue) |
| 486 | |
| 487 | |
| 488 | def run_alphageometry( |