test_trees.py 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101
  1. import pytest
  2. import networkx as nx
  3. from networkx.utils import arbitrary_element, graphs_equal
  4. @pytest.mark.parametrize("prefix_tree_fn", (nx.prefix_tree, nx.prefix_tree_recursive))
  5. def test_basic_prefix_tree(prefix_tree_fn):
  6. # This example is from the Wikipedia article "Trie"
  7. # <https://en.wikipedia.org/wiki/Trie>.
  8. strings = ["a", "to", "tea", "ted", "ten", "i", "in", "inn"]
  9. T = prefix_tree_fn(strings)
  10. root, NIL = 0, -1
  11. def source_label(v):
  12. return T.nodes[v]["source"]
  13. # First, we check that the tree has the expected
  14. # structure. Recall that each node that corresponds to one of
  15. # the input strings has an edge to the NIL node.
  16. #
  17. # Consider the three children at level 1 in the trie.
  18. a, i, t = sorted(T[root], key=source_label)
  19. # Check the 'a' branch.
  20. assert len(T[a]) == 1
  21. nil = arbitrary_element(T[a])
  22. assert len(T[nil]) == 0
  23. # Check the 'i' branch.
  24. assert len(T[i]) == 2
  25. nil, in_ = sorted(T[i], key=source_label)
  26. assert len(T[nil]) == 0
  27. assert len(T[in_]) == 2
  28. nil, inn = sorted(T[in_], key=source_label)
  29. assert len(T[nil]) == 0
  30. assert len(T[inn]) == 1
  31. nil = arbitrary_element(T[inn])
  32. assert len(T[nil]) == 0
  33. # Check the 't' branch.
  34. te, to = sorted(T[t], key=source_label)
  35. assert len(T[to]) == 1
  36. nil = arbitrary_element(T[to])
  37. assert len(T[nil]) == 0
  38. tea, ted, ten = sorted(T[te], key=source_label)
  39. assert len(T[tea]) == 1
  40. assert len(T[ted]) == 1
  41. assert len(T[ten]) == 1
  42. nil = arbitrary_element(T[tea])
  43. assert len(T[nil]) == 0
  44. nil = arbitrary_element(T[ted])
  45. assert len(T[nil]) == 0
  46. nil = arbitrary_element(T[ten])
  47. assert len(T[nil]) == 0
  48. # Next, we check that the "sources" of each of the nodes is the
  49. # rightmost letter in the string corresponding to the path to
  50. # that node.
  51. assert source_label(root) is None
  52. assert source_label(a) == "a"
  53. assert source_label(i) == "i"
  54. assert source_label(t) == "t"
  55. assert source_label(in_) == "n"
  56. assert source_label(inn) == "n"
  57. assert source_label(to) == "o"
  58. assert source_label(te) == "e"
  59. assert source_label(tea) == "a"
  60. assert source_label(ted) == "d"
  61. assert source_label(ten) == "n"
  62. assert source_label(NIL) == "NIL"
  63. @pytest.mark.parametrize(
  64. "strings",
  65. (
  66. ["a", "to", "tea", "ted", "ten", "i", "in", "inn"],
  67. ["ab", "abs", "ad"],
  68. ["ab", "abs", "ad", ""],
  69. ["distant", "disparaging", "distant", "diamond", "ruby"],
  70. ),
  71. )
  72. def test_implementations_consistent(strings):
  73. """Ensure results are consistent between prefix_tree implementations."""
  74. assert graphs_equal(nx.prefix_tree(strings), nx.prefix_tree_recursive(strings))
  75. def test_random_tree():
  76. """Tests that a random tree is in fact a tree."""
  77. T = nx.random_tree(10, seed=1234)
  78. assert nx.is_tree(T)
  79. def test_random_tree_n_zero():
  80. """Tests if n = 0 then the NetworkXPointlessConcept exception is raised."""
  81. with pytest.raises(nx.NetworkXPointlessConcept):
  82. T = nx.random_tree(0, seed=1234)
  83. def test_random_tree_using_generator():
  84. """Tests that creating a ramdom tree with a generator works"""
  85. G = nx.Graph()
  86. T = nx.random_tree(10, seed=1234, create_using=G)
  87. assert nx.is_tree(T)