diff --git a/src/iterators/dfs.py b/src/iterators/dfs.py new file mode 100644 index 0000000..e0cf9a7 --- /dev/null +++ b/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 diff --git a/tests/dfs_test.py b/tests/dfs_test.py new file mode 100644 index 0000000..655180a --- /dev/null +++ b/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)