Coverage for blog/dsa/graph/__init__.py: 71%

1import secrets 

2from typing import Any, Generator 

3 

4from typing_extensions import Self 

5 

6from dsa import queue 

7 

8 

9# start snippet node 

10class Node: 

11 identity: str 

12 color: int | None 

13 neighbors: set[tuple[int, Self]] 

14 

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() 

25 

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) 

31 

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}" 

38 

39 def __hash__(self) -> int: 

40 return int(self.identity, 16) 

41 

42 @classmethod 

43 def _layers( 

44 cls, previous: set[tuple[int, Self]], exclude: set[Self] 

45 ) -> set[tuple[int, Self]]: 

46 

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 

53 

54 out.add((node_distance, node)) 

55 visited.add(node) 

56 

57 return out 

58 

59 def layers(self) -> Generator[set[tuple[int, Self]], None, None]: 

60 exclude = set() 

61 layer = {(0, self)} 

62 

63 while len(layer): 

64 yield layer 

65 exclude |= set(map(lambda item: item[1], layer)) 

66 layer = self._layers(layer, exclude) 

67 

68 # end snippet node 

69 

70 # ----------------------------------------------------------------------- # 

71 # helpers 

72 

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)) 

77 

78 @classmethod 

79 def _from_dict(cls, raw: dict[str, Any]): 

80 

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 

88 

89 @classmethod 

90 def from_dict(cls, raw: dict[str, Any]): 

91 return cls._from_dict(raw.copy()) 

92 

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) 

99 

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 ): 

108 

109 visited.add(self) 

110 

111 out = { 

112 "identity": self.identity, 

113 "color": self.color, 

114 } 

115 if _weight is not None: 

116 out["weight"] = _weight 

117 

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 ] 

124 

125 return out 

126 

127 

128def dijkstra(node_start: Node, node_stop: Node): 

129 

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] 

133 

134 node = node_start 

135 node.color = 0 

136 

137 visited: set[Node] = set() 

138 path = queue.Queue[Node]() 

139 

140 while True: 

141 

142 path.enqueue(node) 

143 if node == node_stop: 

144 break 

145 

146 if not (next_layer := node._layers({(0, node)}, visited)): 

147 return None 

148 

149 _, next_node = min(next_layer, key=lambda item: key(node, *item)) 

150 visited.add(node) 

151 node = next_node 

152 

153 return path