summaryrefslogtreecommitdiff
path: root/networkx/algorithms/traversal/tests/test_beamsearch.py
blob: 8945b418457b0c5c31b956f40a9c4ded2a439d9d (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
"""Unit tests for the beam search functions."""

import networkx as nx


def identity(x):
    return x


class TestBeamSearch:
    """Unit tests for the beam search function."""

    def test_narrow(self):
        """Tests that a narrow beam width may cause an incomplete search."""
        # In this search, we enqueue only the neighbor 3 at the first
        # step, then only the neighbor 2 at the second step. Once at
        # node 2, the search chooses node 3, since it has a higher value
        # that node 1, but node 3 has already been visited, so the
        # search terminates.
        G = nx.cycle_graph(4)
        edges = nx.bfs_beam_edges(G, 0, identity, width=1)
        assert list(edges) == [(0, 3), (3, 2)]

    def test_wide(self):
        G = nx.cycle_graph(4)
        edges = nx.bfs_beam_edges(G, 0, identity, width=2)
        assert list(edges) == [(0, 3), (0, 1), (3, 2)]

    def test_width_none(self):
        G = nx.cycle_graph(4)
        edges = nx.bfs_beam_edges(G, 0, identity, width=None)
        assert list(edges) == [(0, 3), (0, 1), (3, 2)]