Iterators

src/iterators/dfs.py

@@ -0,0 +1,53 @@
+class Graph:
+    def __init__(self, edges: list[tuple[int, int]]):
+        self.edges = edges
+        self.vertices = []
+        for edge in edges:
+            if edge[0] not in self.vertices:
+                self.vertices.append(edge[0])
+            if edge[1] not in self.vertices:
+                self.vertices.append(edge[1])
+
+    def dfs(self) -> list[int]:
+        """Depth-first search algorithm. Returns list of visited vertices."""
+        vertices_state = {
+            "white": self.vertices.copy(),
+            "gray": list(),
+            "black": list(),
+        }
+        visited_vertices_list = []
+
+        def dfs_step(vertex):
+            if vertex in vertices_state["white"]:
+                visited_vertices_list.append(vertex)
+                vertices_state["white"].remove(vertex)
+                vertices_state["gray"].append(vertex)
+
+                for edge in self.edges:
+                    if edge[0] == vertex:
+                        dfs_step(edge[1])
+                    if edge[1] == vertex:
+                        dfs_step(edge[0])
+
+                vertices_state["gray"].remove(vertex)
+                vertices_state["black"].append(vertex)
+
+        for vertex in self.vertices:
+            dfs_step(vertex)
+
+        return visited_vertices_list
+
+    def __iter__(self):
+        return _GraphIterable(self.dfs())
+
+
+class _GraphIterable:
+    def __init__(self, dfs_vertices_list):
+        self.dfs_vertices_list = dfs_vertices_list
+        self.index = 0
+
+    def __next__(self):
+        if self.index < len(self.dfs_vertices_list):
+            self.index += 1
+            return self.dfs_vertices_list[self.index - 1]
+        raise StopIteration

tests/dfs_test.py

@@ -0,0 +1,63 @@
+from iterators.dfs import Graph
+from random import randint
+
+
+def test_method_simple():
+    graph = Graph([(1, 2), (2, 4), (5, 3)])
+    assert graph.dfs() == [1, 2, 4, 5, 3]
+
+
+def test_iterator_simple():
+    graph = Graph([(1, 2), (2, 4), (5, 3)])
+    assert list(graph) == [1, 2, 4, 5, 3]
+
+
+def test_method_zero_elements():
+    graph = Graph([])
+    assert graph.dfs() == []
+
+
+def test_iterator_zero_elements():
+    graph = Graph([])
+    assert list(graph) == []
+
+
+def test_method_edge_to_same_element():
+    graph = Graph([(1, 2), (2, 4), (3, 3)])
+    assert graph.dfs() == [1, 2, 4, 3]
+
+
+def test_iterator_edge_to_same_element():
+    graph = Graph([(1, 2), (2, 4), (3, 3)])
+    assert list(graph) == [1, 2, 4, 3]
+
+
+def test_iterator_possibility_of_parallel_working():
+    graph = Graph([(1, 2), (2, 4)])
+    result = []
+    for i in graph:
+        for j in graph:
+            result.append((i, j))
+    assert result == [
+        (1, 1),
+        (1, 2),
+        (1, 4),
+        (2, 1),
+        (2, 2),
+        (2, 4),
+        (4, 1),
+        (4, 2),
+        (4, 4),
+    ]
+
+
+def test_random_correct_vertices():
+    random_edges = [
+        (randint(-1000, 1000), randint(-1000, 1000)) for _ in range(randint(1, 1000))
+    ]
+    random_edges_vertices_set = set()
+    for edge in random_edges:
+        random_edges_vertices_set.add(edge[0])
+        random_edges_vertices_set.add(edge[1])
+    graph = Graph(random_edges)
+    assert sorted(graph.dfs()) == sorted(random_edges_vertices_set)