1"""
2Ego graph.
3"""
4
5__all__ = ["ego_graph"]
6
7import networkx as nx
8
9
10@nx._dispatchable(preserve_all_attrs=True, returns_graph=True)
11def ego_graph(G, n, radius=1, center=True, undirected=False, distance=None):
12 """Returns induced subgraph of neighbors centered at node n within
13 a given radius.
14
15 Parameters
16 ----------
17 G : graph
18 A NetworkX Graph or DiGraph
19
20 n : node
21 A single node
22
23 radius : number, optional
24 Include all neighbors of distance<=radius from n.
25
26 center : bool, optional
27 If False, do not include center node in graph
28
29 undirected : bool, optional
30 If True use both in- and out-neighbors of directed graphs.
31
32 distance : key, optional
33 Use specified edge data key as distance. For example, setting
34 distance='weight' will use the edge weight to measure the
35 distance from the node n.
36
37 Notes
38 -----
39 For directed graphs D this produces the "out" neighborhood
40 or successors. If you want the neighborhood of predecessors
41 first reverse the graph with D.reverse(). If you want both
42 directions use the keyword argument undirected=True.
43
44 Node, edge, and graph attributes are copied to the returned subgraph.
45 """
46 if undirected:
47 if distance is not None:
48 sp, _ = nx.single_source_dijkstra(
49 G.to_undirected(), n, cutoff=radius, weight=distance
50 )
51 else:
52 sp = dict(
53 nx.single_source_shortest_path_length(
54 G.to_undirected(), n, cutoff=radius
55 )
56 )
57 else:
58 if distance is not None:
59 sp, _ = nx.single_source_dijkstra(G, n, cutoff=radius, weight=distance)
60 else:
61 sp = nx.single_source_shortest_path_length(G, n, cutoff=radius)
62
63 H = G.subgraph(sp).copy()
64 if not center:
65 H.remove_node(n)
66 return H