ioftools / networkxMiCe / networkxmaster / networkx / generators / tests / test_trees.py @ 5cef0f13
History  View  Annotate  Download (2.77 KB)
1 
from nose.tools import assert_equal, assert_true 

2  
3 
import networkx as nx 
4 
from networkx.generators.trees import NIL 
5 
from networkx.utils import arbitrary_element 
6  
7  
8 
class TestPrefixTree(object): 
9 
"""Unit tests for the prefix tree generator function."""

10  
11 
def test_basic(self): 
12 
# This example is from the Wikipedia article "Trie"

13 
# <https://en.wikipedia.org/wiki/Trie>.

14 
strings = ['a', 'to', 'tea', 'ted', 'ten', 'i', 'in', 'inn'] 
15 
T, root = nx.prefix_tree(strings) 
16  
17 
def source_label(v): return T.node[v]['source'] 
18  
19 
# First, we check that the tree has the expected

20 
# structure. Recall that each node that corresponds to one of

21 
# the input strings has an edge to the NIL node.

22 
#

23 
# Consider the three children at level 1 in the trie.

24 
a, i, t = sorted(T[root], key=source_label)

25 
# Check the 'a' branch.

26 
assert_equal(len(T[a]), 1) 
27 
nil = arbitrary_element(T[a]) 
28 
assert_equal(len(T[nil]), 0) 
29 
# Check the 'i' branch.

30 
assert_equal(len(T[i]), 2) 
31 
nil, in_ = sorted(T[i], key=source_label)

32 
assert_equal(len(T[nil]), 0) 
33 
assert_equal(len(T[in_]), 2) 
34 
nil, inn = sorted(T[in_], key=source_label)

35 
assert_equal(len(T[nil]), 0) 
36 
assert_equal(len(T[inn]), 1) 
37 
nil = arbitrary_element(T[inn]) 
38 
assert_equal(len(T[nil]), 0) 
39 
# Check the 't' branch.

40 
te, to = sorted(T[t], key=source_label)

41 
assert_equal(len(T[to]), 1) 
42 
nil = arbitrary_element(T[to]) 
43 
assert_equal(len(T[nil]), 0) 
44 
tea, ted, ten = sorted(T[te], key=source_label)

45 
assert_equal(len(T[tea]), 1) 
46 
assert_equal(len(T[ted]), 1) 
47 
assert_equal(len(T[ten]), 1) 
48 
nil = arbitrary_element(T[tea]) 
49 
assert_equal(len(T[nil]), 0) 
50 
nil = arbitrary_element(T[ted]) 
51 
assert_equal(len(T[nil]), 0) 
52 
nil = arbitrary_element(T[ten]) 
53 
assert_equal(len(T[nil]), 0) 
54  
55 
# Next, we check that the "sources" of each of the nodes is the

56 
# rightmost letter in the string corresponding to the path to

57 
# that node.

58 
assert_equal(source_label(root), None)

59 
assert_equal(source_label(a), 'a')

60 
assert_equal(source_label(i), 'i')

61 
assert_equal(source_label(t), 't')

62 
assert_equal(source_label(in_), 'n')

63 
assert_equal(source_label(inn), 'n')

64 
assert_equal(source_label(to), 'o')

65 
assert_equal(source_label(te), 'e')

66 
assert_equal(source_label(tea), 'a')

67 
assert_equal(source_label(ted), 'd')

68 
assert_equal(source_label(ten), 'n')

69 
assert_equal(source_label(NIL), NIL) 
70  
71  
72 
def test_random_tree(): 
73 
"""Tests that a random tree is in fact a tree."""

74 
T = nx.random_tree(10, seed=1234) 
75 
assert_true(nx.is_tree(T)) 