Coverage for blog/dsa/graph/__init__.py: 71%
78 statements
« prev ^ index » next coverage.py v7.6.12, created at 2025-02-20 16:23 +0000
« prev ^ index » next coverage.py v7.6.12, created at 2025-02-20 16:23 +0000
1import secrets
2from typing import Any, Generator
4from typing_extensions import Self
6from dsa import queue
9# start snippet node
10class Node:
11 identity: str
12 color: int | None
13 neighbors: set[tuple[int, Self]]
15 def __init__(
16 self,
17 color: int | None = None,
18 *,
19 neighbors: set[tuple[int, Self]] | None = None,
20 identity: str | None = None,
21 ):
22 self.identity = identity if identity is not None else secrets.token_hex(4)
23 self.color = color
24 self.neighbors = neighbors if neighbors is not None else set()
26 # NOTE: Ensure symetry. Since this graph is not directed, A a neighbor
27 # of B implies B is a neighbor of A. Think about how the relation
28 # on the graph could be represented as an adjaceny matrix and how
29 # the matrix would be the transpose of itself.
30 self.connect(*self.neighbors)
32 def __str__(self):
33 n = "\n".join(
34 f" - identity: {node.identity}\n distance: {distance}\n color: {node.color}"
35 for distance, node in self.neighbors
36 )
37 return f"identity: {self.identity}\ncolor: {self.color}:\nneighbors:\n{n}"
39 def __hash__(self) -> int:
40 return int(self.identity, 16)
42 @classmethod
43 def _layers(
44 cls, previous: set[tuple[int, Self]], exclude: set[Self]
45 ) -> set[tuple[int, Self]]:
47 out = set()
48 visited = set()
49 for prev_distance, prev in previous:
50 for node_distance, node in prev.neighbors:
51 if node in exclude or node in visited:
52 continue
54 out.add((node_distance, node))
55 visited.add(node)
57 return out
59 def layers(self) -> Generator[set[tuple[int, Self]], None, None]:
60 exclude = set()
61 layer = {(0, self)}
63 while len(layer):
64 yield layer
65 exclude |= set(map(lambda item: item[1], layer))
66 layer = self._layers(layer, exclude)
68 # end snippet node
70 # ----------------------------------------------------------------------- #
71 # helpers
73 def connect(self, *nodes: tuple[int, Self]):
74 self.neighbors |= {(color, node) for color, node in nodes}
75 for color, node in nodes:
76 node.neighbors.add((color, self))
78 @classmethod
79 def _from_dict(cls, raw: dict[str, Any]):
81 neighbors = raw.pop("neighbors", [])
82 node = cls(**raw)
83 neighbors = (
84 (raw_item.pop("weight"), cls._from_dict(raw_item)) for raw_item in neighbors
85 )
86 node.connect(*neighbors)
87 return node
89 @classmethod
90 def from_dict(cls, raw: dict[str, Any]):
91 return cls._from_dict(raw.copy())
93 def to_dict(
94 self,
95 depth: int = 1,
96 exclude: set[Self] | None = None,
97 ):
98 return self._to_dict(set() if exclude is None else exclude.copy(), depth=depth)
100 def _to_dict(
101 self,
102 visited: set[Self],
103 depth: int = 1,
104 *,
105 _weight: int | None = None,
106 _depth: int = 0,
107 ):
109 visited.add(self)
111 out = {
112 "identity": self.identity,
113 "color": self.color,
114 }
115 if _weight is not None:
116 out["weight"] = _weight
118 if _depth < depth:
119 out["neighbors"] = [ # type: ignore[assignment]
120 node._to_dict(visited, depth, _weight=_weight, _depth=_depth + 1)
121 for _weight, node in self.neighbors
122 if node not in visited
123 ]
125 return out
128def dijkstra(node_start: Node, node_stop: Node):
130 def key(node: Node, weight: int, next_node: Node):
131 next_node.color = weight + node.color # type: ignore[operator]
132 return weight + node.color # type: ignore[operator]
134 node = node_start
135 node.color = 0
137 visited: set[Node] = set()
138 path = queue.Queue[Node]()
140 while True:
142 path.enqueue(node)
143 if node == node_stop:
144 break
146 if not (next_layer := node._layers({(0, node)}, visited)):
147 return None
149 _, next_node = min(next_layer, key=lambda item: key(node, *item))
150 visited.add(node)
151 node = next_node
153 return path