Skip to content

Commit 6aa97c7

Browse files
authored
Switch to using heapq (#38)
The docstring of `SearchNode.__lt__` says you're using "heapq". This patch actually makes it so ;-) and also improves efficiency.
1 parent 6cbe744 commit 6aa97c7

File tree

2 files changed

+12
-7
lines changed

2 files changed

+12
-7
lines changed

astar/__init__.py

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from typing import Callable, Dict, Iterable, Union, TypeVar, Generic
66
from math import inf as infinity
77
from operator import attrgetter
8-
import sortedcontainers # type: ignore
8+
import heapq
99

1010
# introduce generic type
1111
T = TypeVar("T")
@@ -48,23 +48,29 @@ def __missing__(self, k) -> SearchNode[T]:
4848

4949
class OpenSet(Generic[SNType]):
5050
def __init__(self) -> None:
51-
self.sortedlist = sortedcontainers.SortedList(key=attrgetter("fscore"))
51+
self.heap: list[SNType] = []
5252

5353
def push(self, item: SNType) -> None:
5454
item.in_openset = True
55-
self.sortedlist.add(item)
55+
heapq.heappush(self.heap, item)
5656

5757
def pop(self) -> SNType:
58-
item = self.sortedlist.pop(0)
58+
item = heapq.heappop(self.heap)
5959
item.in_openset = False
6060
return item
6161

6262
def remove(self, item: SNType) -> None:
63-
self.sortedlist.remove(item)
63+
idx = self.heap.index(item)
6464
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)
6571

6672
def __len__(self) -> int:
67-
return len(self.sortedlist)
73+
return len(self.heap)
6874

6975

7076
################################################################################*

pyproject.toml

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@ packages = [
1919

2020
[tool.poetry.dependencies]
2121
python = ">=3.8,<4.0.0"
22-
sortedcontainers = ">=2.4.0"
2322

2423
[tool.poetry.dev-dependencies]
2524
black = ">=23.3.0"

0 commit comments

Comments
 (0)