Skip to content

Graph Data Structures

Structures for representing networks and complex relationships.

Graph

Bases: Generic[T]

Source code in pure_python_ds/graphs/graph.py
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
class Graph(Generic[T]):
    def __init__(self, directed: bool = False):
        # Updated to store {vertex: {neighbor: weight}}
        self._adj_list: Dict[T, Dict[T, float]] = {}
        self.directed = directed

    def add_vertex(self, vertex: T) -> None:
        if vertex not in self._adj_list:
            self._adj_list[vertex] = {}

    def add_edge(self, v1: T, v2: T, weight: float = 1.0) -> None:
        """Adds a weighted edge. Defaults to 1.0 if not specified."""
        self.add_vertex(v1)
        self.add_vertex(v2)
        self._adj_list[v1][v2] = weight
        if not self.directed:
            self._adj_list[v2][v1] = weight

    def dijkstra(self, start: T) -> Dict[T, float]:
        """
        Greedy algorithm to find the shortest distance from 'start' to all other nodes.
        Returns a dictionary of {vertex: min_distance}.
        """
        if start not in self._adj_list:
            raise ValueError(f"Start node '{start}' not found in graph.")

        # Distances from start to all other nodes are initially infinity
        distances = {vertex: float("inf") for vertex in self._adj_list}
        distances[start] = 0

        # Priority Queue: (distance, tie_breaker, vertex)
        tie_breaker = 0
        pq: List[Tuple[float, int, T]] = [(0.0, tie_breaker, start)]

        while pq:
            current_distance, _, current_vertex = heapq.heappop(pq)

            # If we found a longer path than we already recorded, skip it
            if current_distance > distances[current_vertex]:
                continue

            for neighbor, weight in self._adj_list[current_vertex].items():
                distance = current_distance + weight

                # If this path is shorter, update and push to PQ
                if distance < distances[neighbor]:
                    distances[neighbor] = distance
                    tie_breaker += 1
                    heapq.heappush(pq, (distance, tie_breaker, neighbor))

        return distances

    def kruskal_mst(self) -> List[Tuple[T, T, float]]:
        """
        Finds the Minimum Spanning Tree using the library's DSU module.
        Returns a list of edges (v1, v2, weight) in the MST.
        """
        from pure_python_ds.graphs.disjoint_set import DisjointSet

        edges = []
        # Gather all unique edges
        for u in self._adj_list:
            for v, weight in self._adj_list[u].items():
                if self.directed or (v, u, weight) not in edges:
                    edges.append((u, v, weight))

        # Sort edges by weight
        edges.sort(key=lambda x: x[2])

        dsu = DisjointSet[T](self._adj_list.keys())
        mst = []

        for u, v, weight in edges:
            if dsu.union(u, v):
                mst.append((u, v, weight))

        return mst

    def bellman_ford(self, start: T) -> Dict[T, float]:
        """
        Shortest path algorithm that handles negative weights.
        Returns distances or raises ValueError if a negative cycle is found.
        """
        distances = {vertex: float("inf") for vertex in self._adj_list}
        distances[start] = 0

        # Relax edges |V| - 1 times
        for _ in range(len(self._adj_list) - 1):
            for u in self._adj_list:
                for v, weight in self._adj_list[u].items():
                    if distances[u] + weight < distances[v]:
                        distances[v] = distances[u] + weight

        # Check for negative cycles
        for u in self._adj_list:
            for v, weight in self._adj_list[u].items():
                if distances[u] + weight < distances[v]:
                    raise ValueError("Graph contains a negative weight cycle")

        return distances

    def has_edge(self, u: T, v: T) -> bool:
        """Returns True if there is an edge from u to v."""
        return u in self._adj_list and v in self._adj_list[u]

    def topological_sort(self) -> List[Any]:
        """
        Performs a topological sort on a Directed Acyclic Graph (DAG)
        using Kahn's Algorithm.

        Returns:
            A list of vertices in topologically sorted order.

        Raises:
            ValueError: If the graph contains a cycle (sort is impossible).
        """
        # 1. Initialize all in-degrees to 0
        in_degree = {node: 0 for node in self._adj_list}

        # 2. Calculate actual in-degrees
        for node in self._adj_list:
            for neighbor in self._adj_list[node]:
                if neighbor not in in_degree:
                    in_degree[neighbor] = 0
                in_degree[neighbor] += 1

        # 3. Queue nodes with no incoming dependencies (in-degree == 0)
        queue = deque([node for node in in_degree if in_degree[node] == 0])
        sorted_order = []

        # 4. Process the queue
        while queue:
            current = queue.popleft()
            sorted_order.append(current)

            # Reduce the in-degree of all neighbors by 1
            for neighbor in self._adj_list.get(current, []):
                in_degree[neighbor] -= 1
                if in_degree[neighbor] == 0:
                    queue.append(neighbor)

        # 5. Cycle Detection: If we didn't process every node, there's a loop
        if len(sorted_order) != len(in_degree):
            raise ValueError(
                "Graph contains a cycle; topological sort is not possible."
            )

        return sorted_order

    def bfs(self, start: T) -> List[T]:
        """
        Breadth-First Search. Returns the list of visited nodes in order.
        """
        if start not in self._adj_list:
            raise ValueError(f"Start node '{start}' not found in graph.")

        visited: Set[T] = set()
        queue = deque([start])
        order: List[T] = []

        while queue:
            node = queue.popleft()
            if node not in visited:
                visited.add(node)
                order.append(node)
                for neighbor in self._adj_list.get(node, {}):
                    if neighbor not in visited:
                        queue.append(neighbor)
        return order

    def dfs(self, start: T) -> List[T]:
        """
        Depth-First Search (Iterative). Returns visited nodes in order.
        """
        if start not in self._adj_list:
            raise ValueError(f"Start node '{start}' not found in graph.")

        visited: Set[T] = set()
        stack = [start]
        order: List[T] = []

        while stack:
            node = stack.pop()
            if node not in visited:
                visited.add(node)
                order.append(node)
                # To maintain standard DFS order, we add neighbors in reverse
                neighbors = list(self._adj_list.get(node, {}).keys())
                for neighbor in reversed(neighbors):
                    if neighbor not in visited:
                        stack.append(neighbor)
        return order

    def prims_mst(self) -> List[Tuple[T, T, float]]:
        """
        Prim's Minimum Spanning Tree algorithm.
        Returns a list of edges (u, v, weight) forming the MST.
        """
        if not self._adj_list:
            return []

        start = next(iter(self._adj_list))
        visited = {start}
        mst = []
        pq: List[Tuple[float, int, T, T]] = []
        tie_breaker = 0

        for neighbor, weight in self._adj_list[start].items():
            tie_breaker += 1
            heapq.heappush(pq, (weight, tie_breaker, start, neighbor))

        while pq and len(visited) < len(self._adj_list):
            weight, _, u, v = heapq.heappop(pq)
            if v not in visited:
                visited.add(v)
                mst.append((u, v, weight))
                for next_neighbor, next_weight in self._adj_list[v].items():
                    if next_neighbor not in visited:
                        tie_breaker += 1
                        heapq.heappush(pq, (next_weight, tie_breaker, v, next_neighbor))

        return mst

    def a_star_search(self, start: T, goal: T, heuristic: Callable[[T], float]) -> List[T]:
        """
        A* Search algorithm.
        Returns the shortest path (list of nodes) from start to goal.
        """
        if start not in self._adj_list or goal not in self._adj_list:
            raise ValueError("Start or goal node not found in graph.")

        open_set: List[Tuple[float, int, T]] = [(heuristic(start), 0, start)]
        came_from: Dict[T, T] = {}

        g_score = {node: float("inf") for node in self._adj_list}
        g_score[start] = 0

        f_score = {node: float("inf") for node in self._adj_list}
        f_score[start] = heuristic(start)

        tie_breaker = 0

        while open_set:
            _, _, current = heapq.heappop(open_set)

            if current == goal:
                path = [current]
                while current in came_from:
                    current = came_from[current]
                    path.append(current)
                return path[::-1]

            for neighbor, weight in self._adj_list.get(current, {}).items():
                tentative_g_score = g_score[current] + weight

                if tentative_g_score < g_score[neighbor]:
                    came_from[neighbor] = current
                    g_score[neighbor] = tentative_g_score
                    f_score[neighbor] = tentative_g_score + heuristic(neighbor)

                    tie_breaker += 1
                    heapq.heappush(open_set, (f_score[neighbor], tie_breaker, neighbor))

        return []

A* Search algorithm. Returns the shortest path (list of nodes) from start to goal.

Source code in pure_python_ds/graphs/graph.py
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
def a_star_search(self, start: T, goal: T, heuristic: Callable[[T], float]) -> List[T]:
    """
    A* Search algorithm.
    Returns the shortest path (list of nodes) from start to goal.
    """
    if start not in self._adj_list or goal not in self._adj_list:
        raise ValueError("Start or goal node not found in graph.")

    open_set: List[Tuple[float, int, T]] = [(heuristic(start), 0, start)]
    came_from: Dict[T, T] = {}

    g_score = {node: float("inf") for node in self._adj_list}
    g_score[start] = 0

    f_score = {node: float("inf") for node in self._adj_list}
    f_score[start] = heuristic(start)

    tie_breaker = 0

    while open_set:
        _, _, current = heapq.heappop(open_set)

        if current == goal:
            path = [current]
            while current in came_from:
                current = came_from[current]
                path.append(current)
            return path[::-1]

        for neighbor, weight in self._adj_list.get(current, {}).items():
            tentative_g_score = g_score[current] + weight

            if tentative_g_score < g_score[neighbor]:
                came_from[neighbor] = current
                g_score[neighbor] = tentative_g_score
                f_score[neighbor] = tentative_g_score + heuristic(neighbor)

                tie_breaker += 1
                heapq.heappush(open_set, (f_score[neighbor], tie_breaker, neighbor))

    return []

add_edge(v1, v2, weight=1.0)

Adds a weighted edge. Defaults to 1.0 if not specified.

Source code in pure_python_ds/graphs/graph.py
19
20
21
22
23
24
25
def add_edge(self, v1: T, v2: T, weight: float = 1.0) -> None:
    """Adds a weighted edge. Defaults to 1.0 if not specified."""
    self.add_vertex(v1)
    self.add_vertex(v2)
    self._adj_list[v1][v2] = weight
    if not self.directed:
        self._adj_list[v2][v1] = weight

bellman_ford(start)

Shortest path algorithm that handles negative weights. Returns distances or raises ValueError if a negative cycle is found.

Source code in pure_python_ds/graphs/graph.py
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
def bellman_ford(self, start: T) -> Dict[T, float]:
    """
    Shortest path algorithm that handles negative weights.
    Returns distances or raises ValueError if a negative cycle is found.
    """
    distances = {vertex: float("inf") for vertex in self._adj_list}
    distances[start] = 0

    # Relax edges |V| - 1 times
    for _ in range(len(self._adj_list) - 1):
        for u in self._adj_list:
            for v, weight in self._adj_list[u].items():
                if distances[u] + weight < distances[v]:
                    distances[v] = distances[u] + weight

    # Check for negative cycles
    for u in self._adj_list:
        for v, weight in self._adj_list[u].items():
            if distances[u] + weight < distances[v]:
                raise ValueError("Graph contains a negative weight cycle")

    return distances

bfs(start)

Breadth-First Search. Returns the list of visited nodes in order.

Source code in pure_python_ds/graphs/graph.py
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
def bfs(self, start: T) -> List[T]:
    """
    Breadth-First Search. Returns the list of visited nodes in order.
    """
    if start not in self._adj_list:
        raise ValueError(f"Start node '{start}' not found in graph.")

    visited: Set[T] = set()
    queue = deque([start])
    order: List[T] = []

    while queue:
        node = queue.popleft()
        if node not in visited:
            visited.add(node)
            order.append(node)
            for neighbor in self._adj_list.get(node, {}):
                if neighbor not in visited:
                    queue.append(neighbor)
    return order

dfs(start)

Depth-First Search (Iterative). Returns visited nodes in order.

Source code in pure_python_ds/graphs/graph.py
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
def dfs(self, start: T) -> List[T]:
    """
    Depth-First Search (Iterative). Returns visited nodes in order.
    """
    if start not in self._adj_list:
        raise ValueError(f"Start node '{start}' not found in graph.")

    visited: Set[T] = set()
    stack = [start]
    order: List[T] = []

    while stack:
        node = stack.pop()
        if node not in visited:
            visited.add(node)
            order.append(node)
            # To maintain standard DFS order, we add neighbors in reverse
            neighbors = list(self._adj_list.get(node, {}).keys())
            for neighbor in reversed(neighbors):
                if neighbor not in visited:
                    stack.append(neighbor)
    return order

dijkstra(start)

Greedy algorithm to find the shortest distance from 'start' to all other nodes. Returns a dictionary of {vertex: min_distance}.

Source code in pure_python_ds/graphs/graph.py
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def dijkstra(self, start: T) -> Dict[T, float]:
    """
    Greedy algorithm to find the shortest distance from 'start' to all other nodes.
    Returns a dictionary of {vertex: min_distance}.
    """
    if start not in self._adj_list:
        raise ValueError(f"Start node '{start}' not found in graph.")

    # Distances from start to all other nodes are initially infinity
    distances = {vertex: float("inf") for vertex in self._adj_list}
    distances[start] = 0

    # Priority Queue: (distance, tie_breaker, vertex)
    tie_breaker = 0
    pq: List[Tuple[float, int, T]] = [(0.0, tie_breaker, start)]

    while pq:
        current_distance, _, current_vertex = heapq.heappop(pq)

        # If we found a longer path than we already recorded, skip it
        if current_distance > distances[current_vertex]:
            continue

        for neighbor, weight in self._adj_list[current_vertex].items():
            distance = current_distance + weight

            # If this path is shorter, update and push to PQ
            if distance < distances[neighbor]:
                distances[neighbor] = distance
                tie_breaker += 1
                heapq.heappush(pq, (distance, tie_breaker, neighbor))

    return distances

has_edge(u, v)

Returns True if there is an edge from u to v.

Source code in pure_python_ds/graphs/graph.py
110
111
112
def has_edge(self, u: T, v: T) -> bool:
    """Returns True if there is an edge from u to v."""
    return u in self._adj_list and v in self._adj_list[u]

kruskal_mst()

Finds the Minimum Spanning Tree using the library's DSU module. Returns a list of edges (v1, v2, weight) in the MST.

Source code in pure_python_ds/graphs/graph.py
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
def kruskal_mst(self) -> List[Tuple[T, T, float]]:
    """
    Finds the Minimum Spanning Tree using the library's DSU module.
    Returns a list of edges (v1, v2, weight) in the MST.
    """
    from pure_python_ds.graphs.disjoint_set import DisjointSet

    edges = []
    # Gather all unique edges
    for u in self._adj_list:
        for v, weight in self._adj_list[u].items():
            if self.directed or (v, u, weight) not in edges:
                edges.append((u, v, weight))

    # Sort edges by weight
    edges.sort(key=lambda x: x[2])

    dsu = DisjointSet[T](self._adj_list.keys())
    mst = []

    for u, v, weight in edges:
        if dsu.union(u, v):
            mst.append((u, v, weight))

    return mst

prims_mst()

Prim's Minimum Spanning Tree algorithm. Returns a list of edges (u, v, weight) forming the MST.

Source code in pure_python_ds/graphs/graph.py
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
def prims_mst(self) -> List[Tuple[T, T, float]]:
    """
    Prim's Minimum Spanning Tree algorithm.
    Returns a list of edges (u, v, weight) forming the MST.
    """
    if not self._adj_list:
        return []

    start = next(iter(self._adj_list))
    visited = {start}
    mst = []
    pq: List[Tuple[float, int, T, T]] = []
    tie_breaker = 0

    for neighbor, weight in self._adj_list[start].items():
        tie_breaker += 1
        heapq.heappush(pq, (weight, tie_breaker, start, neighbor))

    while pq and len(visited) < len(self._adj_list):
        weight, _, u, v = heapq.heappop(pq)
        if v not in visited:
            visited.add(v)
            mst.append((u, v, weight))
            for next_neighbor, next_weight in self._adj_list[v].items():
                if next_neighbor not in visited:
                    tie_breaker += 1
                    heapq.heappush(pq, (next_weight, tie_breaker, v, next_neighbor))

    return mst

topological_sort()

Performs a topological sort on a Directed Acyclic Graph (DAG) using Kahn's Algorithm.

Returns:

Type Description
List[Any]

A list of vertices in topologically sorted order.

Raises:

Type Description
ValueError

If the graph contains a cycle (sort is impossible).

Source code in pure_python_ds/graphs/graph.py
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
def topological_sort(self) -> List[Any]:
    """
    Performs a topological sort on a Directed Acyclic Graph (DAG)
    using Kahn's Algorithm.

    Returns:
        A list of vertices in topologically sorted order.

    Raises:
        ValueError: If the graph contains a cycle (sort is impossible).
    """
    # 1. Initialize all in-degrees to 0
    in_degree = {node: 0 for node in self._adj_list}

    # 2. Calculate actual in-degrees
    for node in self._adj_list:
        for neighbor in self._adj_list[node]:
            if neighbor not in in_degree:
                in_degree[neighbor] = 0
            in_degree[neighbor] += 1

    # 3. Queue nodes with no incoming dependencies (in-degree == 0)
    queue = deque([node for node in in_degree if in_degree[node] == 0])
    sorted_order = []

    # 4. Process the queue
    while queue:
        current = queue.popleft()
        sorted_order.append(current)

        # Reduce the in-degree of all neighbors by 1
        for neighbor in self._adj_list.get(current, []):
            in_degree[neighbor] -= 1
            if in_degree[neighbor] == 0:
                queue.append(neighbor)

    # 5. Cycle Detection: If we didn't process every node, there's a loop
    if len(sorted_order) != len(in_degree):
        raise ValueError(
            "Graph contains a cycle; topological sort is not possible."
        )

    return sorted_order

Disjoint Set (Union-Find)

Bases: Generic[T]

A Disjoint Set Union (DSU) or Union-Find data structure.

This data structure keeps track of a set of elements partitioned into a number of disjoint (non-overlapping) subsets. It provides near-constant-time operations to add new sets, merge existing sets, and determine whether elements are in the same set.

Optimizations used: - Path Compression: Flattens the structure of the tree whenever find is used. - Union by Rank: Attaches the shorter tree to the root of the taller tree.

Source code in pure_python_ds/graphs/disjoint_set.py
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
class DisjointSet(Generic[T]):
    """
    A Disjoint Set Union (DSU) or Union-Find data structure.

    This data structure keeps track of a set of elements partitioned into a number
    of disjoint (non-overlapping) subsets. It provides near-constant-time
    operations to add new sets, merge existing sets, and determine whether
    elements are in the same set.

    Optimizations used:
    - Path Compression: Flattens the structure of the tree whenever find is used.
    - Union by Rank: Attaches the shorter tree to the root of the taller tree.
    """

    def __init__(self, items: Optional[Iterable[T]] = None) -> None:
        self._parent: dict[T, T] = {}
        self._rank: dict[T, int] = {}
        self._count: int = 0  # Number of disjoint sets

        if items:
            for item in items:
                self.add(item)

    @property
    def count(self) -> int:
        """Return the number of disjoint sets."""
        return self._count

    def add(self, item: T) -> None:
        """
        Add a new element to the set.

        If the element already exists, this operation does nothing.
        """
        if item not in self._parent:
            self._parent[item] = item
            self._rank[item] = 0
            self._count += 1

    def find(self, item: T) -> T:
        """
        Find the representative of the set containing the item.

        Applies path compression.
        """
        if item not in self._parent:
            raise KeyError(f"Item {item!r} not found in DisjointSet")

        if self._parent[item] != item:
            self._parent[item] = self.find(self._parent[item])
        return self._parent[item]

    def union(self, item1: T, item2: T) -> bool:
        """
        Union the sets containing item1 and item2.

        Returns True if a merge happened, False if they were already in the same set.
        Applies union by rank.
        """
        root1 = self.find(item1)
        root2 = self.find(item2)

        if root1 != root2:
            # Union by rank
            if self._rank[root1] < self._rank[root2]:
                self._parent[root1] = root2
            elif self._rank[root1] > self._rank[root2]:
                self._parent[root2] = root1
            else:
                self._parent[root2] = root1
                self._rank[root1] += 1

            self._count -= 1
            return True

        return False

    def connected(self, item1: T, item2: T) -> bool:
        """Check if two items are in the same set."""
        return self.find(item1) == self.find(item2)

    def __contains__(self, item: T) -> bool:
        """Check if an item is in the Disjoint Set."""
        return item in self._parent

    def __len__(self) -> int:
        """Return the number of elements in the Disjoint Set."""
        return len(self._parent)

count property

Return the number of disjoint sets.

__contains__(item)

Check if an item is in the Disjoint Set.

Source code in pure_python_ds/graphs/disjoint_set.py
92
93
94
def __contains__(self, item: T) -> bool:
    """Check if an item is in the Disjoint Set."""
    return item in self._parent

__len__()

Return the number of elements in the Disjoint Set.

Source code in pure_python_ds/graphs/disjoint_set.py
96
97
98
def __len__(self) -> int:
    """Return the number of elements in the Disjoint Set."""
    return len(self._parent)

add(item)

Add a new element to the set.

If the element already exists, this operation does nothing.

Source code in pure_python_ds/graphs/disjoint_set.py
39
40
41
42
43
44
45
46
47
48
def add(self, item: T) -> None:
    """
    Add a new element to the set.

    If the element already exists, this operation does nothing.
    """
    if item not in self._parent:
        self._parent[item] = item
        self._rank[item] = 0
        self._count += 1

connected(item1, item2)

Check if two items are in the same set.

Source code in pure_python_ds/graphs/disjoint_set.py
88
89
90
def connected(self, item1: T, item2: T) -> bool:
    """Check if two items are in the same set."""
    return self.find(item1) == self.find(item2)

find(item)

Find the representative of the set containing the item.

Applies path compression.

Source code in pure_python_ds/graphs/disjoint_set.py
50
51
52
53
54
55
56
57
58
59
60
61
def find(self, item: T) -> T:
    """
    Find the representative of the set containing the item.

    Applies path compression.
    """
    if item not in self._parent:
        raise KeyError(f"Item {item!r} not found in DisjointSet")

    if self._parent[item] != item:
        self._parent[item] = self.find(self._parent[item])
    return self._parent[item]

union(item1, item2)

Union the sets containing item1 and item2.

Returns True if a merge happened, False if they were already in the same set. Applies union by rank.

Source code in pure_python_ds/graphs/disjoint_set.py
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
def union(self, item1: T, item2: T) -> bool:
    """
    Union the sets containing item1 and item2.

    Returns True if a merge happened, False if they were already in the same set.
    Applies union by rank.
    """
    root1 = self.find(item1)
    root2 = self.find(item2)

    if root1 != root2:
        # Union by rank
        if self._rank[root1] < self._rank[root2]:
            self._parent[root1] = root2
        elif self._rank[root1] > self._rank[root2]:
            self._parent[root2] = root1
        else:
            self._parent[root2] = root1
            self._rank[root1] += 1

        self._count -= 1
        return True

    return False