@@ -11,16 +11,25 @@ def _get_cutoff(n, cutoff):
1111 return cutoff + 1 # Inclusive
1212
1313
14- def _bfs_plain (G , source = None , target = None , * , index = None , cutoff = None ):
14+ # Push-pull optimization is possible, but annoying to implement
15+ def _bfs_plain (
16+ G , source = None , target = None , * , index = None , cutoff = None , transpose = False , name = "bfs_plain"
17+ ):
1518 if source is not None :
19+ if source not in G ._key_to_id :
20+ raise KeyError (f"The node { source } )
1621 index = G ._key_to_id [source ]
1722 if target is not None :
23+ if target not in G ._key_to_id :
24+ raise KeyError (f"The node { target } )
1825 dst_id = G ._key_to_id [target ]
1926 else :
2027 dst_id = None
2128 A = G .get_property ("offdiag" )
29+ if transpose and G .is_directed ():
30+ A = A .T # TODO: should we use "AT" instead?
2231 n = A .nrows
23- v = Vector (bool , n , name = "bfs_plain" )
32+ v = Vector (bool , n , name = name )
2433 q = Vector (bool , n , name = "q" )
2534 v [index ] = True
2635 q [index ] = True
@@ -30,16 +39,22 @@ def _bfs_plain(G, source=None, target=None, *, index=None, cutoff=None):
3039 q (~ v .S , replace ) << any_pair_bool (q @ A )
3140 if q .nvals == 0 :
3241 break
42+ v (q .S ) << True
3343 if dst_id is not None and dst_id in q :
3444 break
35- v (q .S ) << True
3645 return v
3746
3847
39- def _bfs_level (G , source , cutoff = None , * , transpose = False , dtype = int ):
48+ def _bfs_level (G , source , target = None , * , cutoff = None , transpose = False , dtype = int ):
4049 if dtype == bool :
4150 dtype = int
4251 index = G ._key_to_id [source ]
52+ if target is not None :
53+ if target not in G ._key_to_id :
54+ raise KeyError (f"The node { target } )
55+ dst_id = G ._key_to_id [target ]
56+ else :
57+ dst_id = None
4358 A = G .get_property ("offdiag" )
4459 if transpose and G .is_directed ():
4560 A = A .T # TODO: should we use "AT" instead?
@@ -55,10 +70,12 @@ def _bfs_level(G, source, cutoff=None, *, transpose=False, dtype=int):
5570 if q .nvals == 0 :
5671 break
5772 v (q .S ) << i
73+ if dst_id is not None and dst_id in q :
74+ break
5875 return v
5976
6077
61- def _bfs_levels (G , nodes , cutoff = None , * , dtype = int ):
78+ def _bfs_levels (G , nodes , * , cutoff = None , dtype = int ):
6279 if dtype == bool :
6380 dtype = int
6481 A = G .get_property ("offdiag" )
@@ -90,7 +107,7 @@ def _bfs_levels(G, nodes, cutoff=None, *, dtype=int):
90107 return D
91108
92109
93- def _bfs_parent (G , source , cutoff = None , * , target = None , transpose = False , dtype = int ):
110+ def _bfs_parent (G , source , target = None , * , cutoff = None , transpose = False , dtype = int ):
94111 if dtype == bool :
95112 dtype = int
96113 index = G ._key_to_id [source ]
0 commit comments