Add function to find densest subgraph

releasenotes/notes/densest_subgraph-1b068f69f80facd4.yaml

@@ -0,0 +1,31 @@
+---
+features:
+ - |
+ Added a new function, :func:`~.densest_subgraph_of_size`, which is used to return a
+ subgraph of given size that has the highest degree of connecitivity between the nodes.
+ For example, if you wanted to find the subgraph of 5 nodes in a 19 node heavy hexagon
+ graph:
+
+ .. jupyter-execute::
+
+ import rustworkx as rx
+ from rustworkx.visualization import mpl_draw
+
+ graph = rx.generators.hexagonal_lattice_graph(4, 5)
+
+ subgraph, node_map = rx.densest_subgraph_of_size(graph, 5)
+ subgraph_edge_set = set(subgraph.edge_list())
+ node_colors = []
+ for node in graph.node_indices():
+ if node in node_map:
+ node_colors.append('red')
+ else:
+ node_colors.append('blue')
+ graph[node] = node
+ edge_colors = []
+ for edge in graph.edge_list():
+ if edge[0] in node_map and edge[1] in node_map:
+ edge_colors.append('red')
+ else:
+ edge_colors.append('blue')
+ mpl_draw(graph, with_labels=True, node_color=node_colors, edge_color=edge_colors, labels=str)

rustworkx-core/src/dense_subgraph.rs

@@ -0,0 +1,212 @@
+// Licensed under the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License. You may obtain
+// a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+// License for the specific language governing permissions and limitations
+// under the License.
+
+use hashbrown::{HashMap, HashSet};
+use std::hash::Hash;
+
+use petgraph::prelude::*;
+use petgraph::visit::{
+ EdgeCount, GraphProp, IntoEdgeReferences, IntoNeighbors, IntoNodeIdentifiers, NodeCount,
+ Visitable,
+};
+
+use rayon::prelude::*;
+
+struct SubsetResult<N> {
+ pub count: usize,
+ pub error: f64,
+ pub map: Vec<N>,
+}
+
+/// Find the most densely connected k-subgraph
+///
+/// This function will return the node indices of the subgraph of `num_nodes` that is the
+/// most densely connected.
+///
+/// This method does not provide any guarantees on the approximation as it
+/// does a naive search using BFS traversal.
+///
+/// # Arguments
+///
+/// * `graph` - The graph to find densest subgraph in.
+/// * `num_nodes` - The number of nodes in the subgraph to find
+/// * `edge_weight_callback` - An optional callable that if specified will be
+/// passed the node indices of each edge in the graph and it is expected to
+/// return a float value. If specified the lowest avg weight for edges in
+/// a found subgraph will be a criteria for selection in addition to the
+/// connectivity of the subgraph.
+/// * `node_weight_callback` - An optional callable that if specified will be
+/// passed the node indices of each node in the graph and it is expected to
+/// return a float value. If specified the lowest avg weight for node of
+/// a found subgraph will be a criteria for selection in addition to the
+/// connectivity of the subgraph.
+///
+/// # Example:
+///
+/// ```rust
+/// use std::convert::Infallible;
+/// use rustworkx_core::petgraph::stable_graph::{StableDiGraph, NodeIndex};
+/// use rustworkx_core::petgraph::visit::IntoEdgeReferences;
+/// use rustworkx_core::generators::grid_graph;
+/// use rustworkx_core::dense_subgraph::densest_subgraph;
+///
+/// type EdgeWeightType = Box<dyn FnMut(<&StableDiGraph<(), ()> as IntoEdgeReferences>::EdgeRef) -> Result<f64, Infallible>>;
+/// type NodeWeightType = Box<dyn FnMut(NodeIndex) -> Result<f64, Infallible>>;
+///
+/// let graph: StableDiGraph<(), ()> = grid_graph(
+/// Some(10),
+/// Some(10),
+/// None,
+/// || {()},
+/// || {()},
+/// false
+/// ).unwrap();
+/// let subgraph_nodes = densest_subgraph(&graph, 10, None::<EdgeWeightType>, None::<NodeWeightType>).unwrap();
+///
+/// let expected = vec![
+/// NodeIndex::new(7), NodeIndex::new(8), NodeIndex::new(17), NodeIndex::new(9),
+/// NodeIndex::new(18), NodeIndex::new(27), NodeIndex::new(19), NodeIndex::new(28),
+/// NodeIndex::new(37), NodeIndex::new(29)
+/// ];
+///
+/// assert_eq!(subgraph_nodes, expected);
+/// ```
+pub fn densest_subgraph<G, H, F, E>(
+ graph: G,
+ num_nodes: usize,
+ edge_weight_callback: Option<H>,
+ node_weight_callback: Option<F>,
+) -> Result<Vec<G::NodeId>, E>
+where
+ G: IntoNodeIdentifiers
+ + IntoEdgeReferences
+ + EdgeCount
+ + GraphProp
+ + NodeCount
+ + IntoNeighbors
+ + Visitable
+ + Sync,
+ G::NodeId: Eq + Hash + Send + Sync,
+ F: FnMut(G::NodeId) -> Result<f64, E>,
+ H: FnMut(G::EdgeRef) -> Result<f64, E>,
+{
+ let node_indices: Vec<G::NodeId> = graph.node_identifiers().collect();
+ let mut edge_weight_map: Option<HashMap<[G::NodeId; 2], f64>> = None;
+ let mut node_weight_map: Option<HashMap<G::NodeId, f64>> = None;
+
+ if edge_weight_callback.is_some() {
+ let mut inner_weight_map: HashMap<[G::NodeId; 2], f64> =
+ HashMap::with_capacity(graph.edge_count());
+ let mut callback = edge_weight_callback.unwrap();
+ for edge in graph.edge_references() {
+ let source = edge.source();
+ let target = edge.target();
+ let weight = callback(edge)?;
+ inner_weight_map.insert([source, target], weight);
+ if !graph.is_directed() {
+ inner_weight_map.insert([target, source], weight);
+ }
+ }
+ edge_weight_map = Some(inner_weight_map);
+ }
+ let mut avg_node_error: f64 = 0.;
+ if node_weight_callback.is_some() {
+ let mut callback = node_weight_callback.unwrap();
+ let mut inner_weight_map: HashMap<G::NodeId, f64> =
+ HashMap::with_capacity(graph.node_count());
+ for node in graph.node_identifiers() {
+ let weight = callback(node)?;
+ avg_node_error += weight;
+ inner_weight_map.insert(node, weight);
+ }
+ avg_node_error /= graph.node_count() as f64;
+ node_weight_map = Some(inner_weight_map);
+ }
+ let reduce_identity_fn = || -> SubsetResult<G::NodeId> {
+ SubsetResult {
+ count: 0,
+ map: Vec::new(),
+ error: f64::INFINITY,
+ }
+ };
+
+ let reduce_fn =
+ |best: SubsetResult<G::NodeId>, curr: SubsetResult<G::NodeId>| -> SubsetResult<G::NodeId> {
+ if edge_weight_map.is_some() || node_weight_map.is_some() {
+ if curr.count >= best.count && curr.error <= best.error {
+ curr
+ } else {
+ best
+ }
+ } else if curr.count > best.count {
+ curr
+ } else {
+ best
+ }
+ };
+
+ let best_result = node_indices
+ .into_par_iter()
+ .filter_map(|index| {
+ let mut subgraph: Vec<[G::NodeId; 2]> = Vec::with_capacity(num_nodes);
+ let mut bfs = Bfs::new(&graph, index);
+ let mut bfs_vec: Vec<G::NodeId> = Vec::with_capacity(num_nodes);
+ let mut bfs_set: HashSet<G::NodeId> = HashSet::with_capacity(num_nodes);
+
+ let mut count = 0;
+ while let Some(node) = bfs.next(&graph) {
+ bfs_vec.push(node);
+ bfs_set.insert(node);
+ count += 1;
+ if count >= num_nodes {
+ break;
+ }
+ }
+ if bfs_vec.len() < num_nodes {
+ return None;
+ }
+ let mut connection_count = 0;
+ for node in &bfs_vec {
+ for nbr in graph.neighbors(*node).filter(|j| bfs_set.contains(j)) {
+ connection_count += 1;
+ subgraph.push([*node, nbr]);
+ }
+ }
+ let mut error = match &edge_weight_map {
+ Some(map) => {
+ subgraph.iter().map(|edge| map[edge]).sum::<f64>() / subgraph.len() as f64
+ }
+ None => 0.,
+ };
+ error *= match &node_weight_map {
+ Some(map) => {
+ let subgraph_node_error_avg =
+ bfs_vec.iter().map(|node| map[node]).sum::<f64>() / num_nodes as f64;
+ let node_error_diff = subgraph_node_error_avg - avg_node_error;
+ if node_error_diff > 0. {
+ num_nodes as f64 * node_error_diff
+ } else {
+ 1.
+ }
+ }
+ None => 1.,
+ };
+
+ Some(SubsetResult {
+ count: connection_count,
+ error,
+ map: bfs_vec,
+ })
+ })
+ .reduce(reduce_identity_fn, reduce_fn);
+ Ok(best_result.map)
+}

rustworkx-core/src/lib.rs

@@ -109,6 +109,7 @@ pub mod planar;
 pub mod shortest_path;
 pub mod traversal;
 // These modules define additional data structures
+pub mod dense_subgraph;
 pub mod dictmap;
 pub mod distancemap;
 mod min_scored;

rustworkx/__init__.py

@@ -1873,6 +1873,33 @@ def longest_simple_path(graph):
  found in the graph. If the graph is empty ``None`` will be returned instead.
  :rtype: NodeIndices
  """
+ raise TypeError("Invalid Input Type %s for graph" % type(graph))
+
+
+@_rustworkx_dispatch
+def densest_subgraph_of_size(
+ graph, num_nodes, /, edge_weight_callback=None, node_weight_callback=None
+):
+ """Find a connected and dense subgraph of a given size in a graph.
+
+ This method does not provide any guarantees on the approximation of the optimal solution as it
+ does a naive search using BFS traversal.
+
+ :param graph: The graph to find the densest subgraph in. This can be a
+ :class:`~retworkx.PyGraph` or a :class:`~retworkx.PyDiGraph`.
+ :param int num_nodes: The number of nodes in the subgraph to find
+ :param func weight_callback: An optional callable that if specified will be
+ passed the node indices of each edge in the graph and it is expected to
+ return a float value. If specified the lowest avg weight for edges in
+ a found subgraph will be a criteria for selection in addition to the
+ connectivity of the subgraph.
+ :returns: A tuple of the subgraph found and a :class:`~.NodeMap` of the
+ mapping of node indices in the input ``graph`` to the index in the
+ output subgraph.
+
+ :rtype: (subgraph, node_map)
+ """
+ raise TypeError("Invalid Input Type %s for graph" % type(graph))
 
 
 @_rustworkx_dispatch

rustworkx/__init__.pyi

@@ -232,6 +232,8 @@ from .rustworkx import steiner_tree as steiner_tree
 from .rustworkx import metric_closure as metric_closure
 from .rustworkx import digraph_union as digraph_union
 from .rustworkx import graph_union as graph_union
+from .rustworkx import digraph_densest_subgraph_of_size as digraph_densest_subgraph_of_size
+from .rustworkx import graph_densest_subgraph_of_size as graph_densest_subgraph_of_size
 from .rustworkx import NodeIndices as NodeIndices
 from .rustworkx import PathLengthMapping as PathLengthMapping
 from .rustworkx import PathMapping as PathMapping
@@ -602,3 +604,10 @@ def longest_simple_path(graph: PyGraph[_S, _T] | PyDiGraph[_S, _T]) -> NodeIndic
 def isolates(graph: PyGraph[_S, _T] | PyDiGraph[_S, _T]) -> NodeIndices: ...
 def two_color(graph: PyGraph[_S, _T] | PyDiGraph[_S, _T]) -> dict[int, int]: ...
 def is_bipartite(graph: PyGraph[_S, _T] | PyDiGraph[_S, _T]) -> bool: ...
+def densest_subgraph_of_size(
+ graph: PyGraph[_S, _T] | PyDiGraph[_S, _T],
+ num_nodes: int,
+ /,
+ edge_weight_callback: Callable[[_T], float] | None = ...,
+ node_weight_callback: Callable[[_S], float] | None = ...,
+) -> tuple[PyGraph[_S, _T], NodeMap]: ...

rustworkx/rustworkx.pyi

@@ -1012,6 +1012,23 @@ def graph_union(
  merge_edges: bool = ...,
 ) -> PyGraph[_S, _T]: ...
 
+# Densest Subgraph
+
+def graph_densest_subgraph_of_size(
+ graph: PyGraph[_S, _T],
+ num_nodes: int,
+ /,
+ edge_weight_callback: Callable[[_T], float] | None = ...,
+ node_weight_callback: Callable[[_T], float] | None = ...,
+) -> tuple[PyGraph[_S, _T], NodeMap]: ...
+def digraph_densest_subgraph_of_size(
+ graph: PyDiGraph[_S, _T],
+ num_nodes: int,
+ /,
+ edge_weight_callback: Callable[[_T], float] | None = ...,
+ node_weight_callback: Callable[[_T], float] | None = ...,
+) -> tuple[PyGraph[_S, _T], NodeMap]: ...
+
 # Iterators
 
 _T_co = TypeVar("_T_co", covariant=True)