Skip to content

Commit 649bd23

Browse files
committed
Adding generator for creating graphs for testing.
1 parent 6c70286 commit 649bd23

File tree

2 files changed

+64
-0
lines changed

2 files changed

+64
-0
lines changed

graph/generate.py

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
import itertools
2+
from graph.core import Graph
3+
4+
5+
def binary_tree(levels):
6+
"""
7+
Generates a binary tree with the given number of levels.
8+
"""
9+
if not isinstance(levels, int):
10+
raise TypeError(f"Expected int, not {type(levels)}")
11+
g = Graph()
12+
for i in range(levels):
13+
g.add_edge(i, 2 * i + 1)
14+
g.add_edge(i, 2 * i + 2)
15+
return g
16+
17+
18+
def grid(length, width, bidirectional=False):
19+
"""
20+
Generates a grid with the given length and width.
21+
"""
22+
if not isinstance(length, int):
23+
raise TypeError(f"Expected int, not {type(length)}")
24+
if not isinstance(width, int):
25+
raise TypeError(f"Expected int, not {type(width)}")
26+
g = Graph()
27+
node_index = {}
28+
c = itertools.count(start=1)
29+
for i in range(length): # i is the row
30+
for j in range(width): # j is the column
31+
node_index[(i, j)] = next(c)
32+
if i > 0:
33+
a, b = node_index[(i, j)], node_index[(i - 1, j)]
34+
g.add_edge(b, a, bidirectional=bidirectional)
35+
if j > 0:
36+
a, b = node_index[(i, j)], node_index[(i, j - 1)]
37+
g.add_edge(b, a, bidirectional=bidirectional)
38+
return g

tests/test_generate.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
from graph.core import Graph
2+
from graph.generate import binary_tree, grid
3+
4+
from tests.test_graph import graph3x3, graph4x4, graph5x5
5+
6+
7+
def test_generate_grid_3x3():
8+
g = grid(3, 3)
9+
expected = graph3x3()
10+
assert g == expected
11+
12+
13+
def test_generate_grid_4x4_bidirectional():
14+
g = grid(4, 4, bidirectional=True)
15+
expected = graph4x4()
16+
assert g == expected
17+
18+
def test_generate_grid_5x5_bidirectional():
19+
g = grid(5, 5, bidirectional=True)
20+
expected = graph5x5()
21+
assert g == expected
22+
23+
def test_generate_binary_tree_3_levels():
24+
g = binary_tree(3)
25+
expected = Graph(from_list=[[0, 1], [0, 2], [1, 3], [1, 4], [2, 5], [2, 6]])
26+
assert g == expected

0 commit comments

Comments
 (0)