|
5 | 5 | from typing import Callable, Dict, Iterable, Union, TypeVar, Generic
|
6 | 6 | from math import inf as infinity
|
7 | 7 | from operator import attrgetter
|
8 |
| -import sortedcontainers # type: ignore |
| 8 | +import heapq |
9 | 9 |
|
10 | 10 | # introduce generic type
|
11 | 11 | T = TypeVar("T")
|
@@ -48,23 +48,29 @@ def __missing__(self, k) -> SearchNode[T]:
|
48 | 48 |
|
49 | 49 | class OpenSet(Generic[SNType]):
|
50 | 50 | def __init__(self) -> None:
|
51 |
| - self.sortedlist = sortedcontainers.SortedList(key=attrgetter("fscore")) |
| 51 | + self.heap: list[SNType] = [] |
52 | 52 |
|
53 | 53 | def push(self, item: SNType) -> None:
|
54 | 54 | item.in_openset = True
|
55 |
| - self.sortedlist.add(item) |
| 55 | + heapq.heappush(self.heap, item) |
56 | 56 |
|
57 | 57 | def pop(self) -> SNType:
|
58 |
| - item = self.sortedlist.pop(0) |
| 58 | + item = heapq.heappop(self.heap) |
59 | 59 | item.in_openset = False
|
60 | 60 | return item
|
61 | 61 |
|
62 | 62 | def remove(self, item: SNType) -> None:
|
63 |
| - self.sortedlist.remove(item) |
| 63 | + idx = self.heap.index(item) |
64 | 64 | item.in_openset = False
|
| 65 | + item = self.heap.pop() |
| 66 | + if idx < len(self.heap): |
| 67 | + self.heap[idx] = item |
| 68 | + # Fix heap invariants |
| 69 | + heapq._siftup(self.heap, idx) |
| 70 | + heapq._siftdown(self.heap, 0, idx) |
65 | 71 |
|
66 | 72 | def __len__(self) -> int:
|
67 |
| - return len(self.sortedlist) |
| 73 | + return len(self.heap) |
68 | 74 |
|
69 | 75 |
|
70 | 76 | ################################################################################*
|
|
0 commit comments