use petgraph::algo::{bellman_ford, min_spanning_tree};
use petgraph::data::FromElements;
use petgraph::prelude::*;
use petgraph::Graph;
use std::collections::HashMap;
fn convert_shortest_paths(
all_paths: &[bellman_ford::Paths<NodeIndex, f64>],
nodes: &[NodeIndex],
) -> Vec<Vec<(usize, Vec<usize>)>> {
let mut my_paths: Vec<Vec<(usize, Vec<usize>)>> = Vec::new();
for n1 in nodes.iter() {
my_paths.push(Vec::new());
for n2 in nodes.iter() {
let path = get_shortest_path(all_paths, n1, n2);
my_paths.last_mut().unwrap().push((path.len() - 1, path));
}
}
my_paths
}
fn get_shortest_path(
all_paths: &[bellman_ford::Paths<NodeIndex, f64>],
n1: &NodeIndex,
n2: &NodeIndex,
) -> Vec<usize> {
let path = &all_paths[n1.index()];
let mut spath: Vec<usize> = Vec::new();
spath.push(n2.index());
let mut v = path.predecessors[n2.index()];
while let Some(ni) = v {
spath.push(ni.index());
v = path.predecessors[ni.index()];
}
spath
}
fn build_closure_mst(
all_paths: &[Vec<(usize, Vec<usize>)>],
terminals: &[usize],
) -> Graph<(), f64, Undirected> {
let mut g = Graph::new_undirected();
let nodes: Vec<_> = (0..terminals.len()).map(|_| g.add_node(())).collect();
for (i1, n1) in nodes.iter().enumerate() {
for (i2, n2) in nodes.iter().enumerate() {
if i1 != i2 {
g.add_edge(*n1, *n2, all_paths[terminals[i1]][terminals[i2]].0 as f64);
}
}
}
let mst: Graph<(), f64, Undirected> = Graph::from_elements(min_spanning_tree(&g));
mst
}
#[derive(Debug, Clone)]
pub struct SteinerTree {
pub graph: UnGraph<(), i32, u32>,
pub mapping: HashMap<usize, NodeIndex>,
pub inverse_mapping: HashMap<NodeIndex, usize>,
pub terminals: Vec<usize>,
}
impl SteinerTree {
pub fn len(&self) -> usize {
self.graph.node_count()
}
pub fn is_empty(&self) -> bool {
self.graph.node_count() == 0
}
pub fn cnot_cost(&self) -> usize {
2 * self.len() - self.terminals.len() - 1
}
pub fn contains(&self, node: usize) -> bool {
self.mapping.contains_key(&node)
}
pub fn is_leaf(&self, node: usize) -> bool {
let node_index = self.mapping.get(&node).unwrap();
self.graph.neighbors_undirected(*node_index).count() == 1
}
pub fn neighbors(&self, node: usize) -> Vec<usize> {
let mut neighs = Vec::new();
for nei in self.graph.neighbors_undirected(self.mapping[&node]) {
neighs.push(self.inverse_mapping[&nei]);
}
neighs
}
pub fn nodes(&self) -> Vec<usize> {
self.mapping.keys().copied().collect()
}
pub fn add_node(&mut self, node: usize) {
let node_index = self.graph.add_node(());
self.mapping.insert(node, node_index);
self.inverse_mapping.insert(node_index, node);
}
pub fn add_edge(&mut self, n1: usize, n2: usize) {
self.graph.add_edge(self.mapping[&n1], self.mapping[&n2], 1);
}
pub fn remove_node(&mut self, node: usize) {
if !self.contains(node) {
panic!("Node {} is not in the tree", node);
}
let last_index: NodeIndex<u32> = NodeIndex::new(self.graph.node_count() - 1);
let node_index = *self.mapping.get(&node).unwrap();
let i = *self.inverse_mapping.get(&last_index).unwrap();
if i == node {
self.graph.remove_node(node_index);
self.mapping.remove(&i);
self.inverse_mapping.remove(&node_index);
return;
}
self.inverse_mapping.remove(&last_index);
self.mapping.remove(&node);
self.mapping.insert(i, node_index);
self.inverse_mapping.insert(node_index, i);
self.graph.remove_node(node_index);
}
pub fn update_tree(&mut self, new_support: &[usize], qbits: &[usize; 2]) {
self.terminals = new_support.to_owned();
if self.contains(qbits[0]) && self.contains(qbits[1]) {
if new_support.contains(&qbits[0]) && new_support.contains(&qbits[1]) {
return;
}
if !new_support.contains(&qbits[0]) && self.is_leaf(qbits[0]) {
self.remove_node(qbits[0]);
}
if !new_support.contains(&qbits[1]) && self.is_leaf(qbits[1]) {
self.remove_node(qbits[1]);
}
return;
}
if !self.contains(qbits[0]) && new_support.contains(&qbits[0]) {
self.add_node(qbits[0]);
self.add_edge(qbits[0], qbits[1]);
return;
}
if !self.contains(qbits[1]) && new_support.contains(&qbits[1]) {
self.add_node(qbits[1]);
self.add_edge(qbits[0], qbits[1]);
}
}
pub fn prune_non_terminal_leaves(&mut self) {
loop {
let mut popped = false;
for node in self.graph.node_indices() {
if !self.terminals.contains(&self.inverse_mapping[&node])
&& self.graph.neighbors(node).count() == 1
{
self.remove_node(self.inverse_mapping[&node]);
popped = true;
break;
}
}
if !popped {
break;
}
}
}
}
fn build_steiner_from_mst(
closure_mst: &Graph<(), f64, Undirected>,
terminals: &[usize],
all_paths: &[Vec<(usize, Vec<usize>)>],
) -> SteinerTree {
let mut g = Graph::new_undirected();
let mut nodes = HashMap::new();
let mut inverse_map = HashMap::new();
for i in terminals.iter() {
let nodei = g.add_node(());
nodes.insert(*i, nodei);
inverse_map.insert(nodei, *i);
}
for edge in closure_mst.raw_edges() {
let t1 = terminals[edge.source().index()];
let t2 = terminals[edge.target().index()];
let (_, path) = &all_paths[t1][t2];
for i in path.iter() {
if !nodes.contains_key(i) {
let nodei = g.add_node(());
nodes.insert(*i, nodei);
inverse_map.insert(nodei, *i);
}
}
let mut left = path[0];
for next_v in path.iter().skip(1) {
g.add_edge(nodes[&left], nodes[next_v], 1);
left = *next_v;
}
}
let mst: Graph<(), i32, Undirected> = Graph::from_elements(min_spanning_tree(&g));
let mut as_st = SteinerTree {
graph: mst,
mapping: nodes,
inverse_mapping: inverse_map,
terminals: terminals.to_owned(),
};
as_st.prune_non_terminal_leaves();
as_st
}
fn get_steiner_tree(conv_paths: &[Vec<(usize, Vec<usize>)>], terminals: &[usize]) -> SteinerTree {
let closure = build_closure_mst(conv_paths, terminals);
build_steiner_from_mst(&closure, terminals, conv_paths)
}
#[derive(Debug, Clone)]
pub struct HardwareGraph {
pub graph: UnGraph<(), f64, u32>,
shortest_paths: Vec<Vec<(usize, Vec<usize>)>>,
}
impl HardwareGraph {
pub fn from_couplings(couplings: &[(usize, usize)]) -> Self {
let mut graph = UnGraph::new_undirected();
let n = couplings.iter().map(|c| c.0.max(c.1)).max().unwrap() + 1;
let nodes: Vec<_> = (0..n).map(|_| graph.add_node(())).collect();
graph.extend_with_edges(couplings.iter().map(|c| (nodes[c.0], nodes[c.1], 1.)));
let all_paths: Vec<bellman_ford::Paths<NodeIndex, f64>> = nodes
.iter()
.map(|node_index| bellman_ford(&graph, *node_index).unwrap())
.collect();
let shortest_paths = convert_shortest_paths(&all_paths, &nodes);
Self {
graph,
shortest_paths,
}
}
pub fn len(&self) -> usize {
self.graph.node_count()
}
pub fn is_empty(&self) -> bool {
self.graph.node_count() == 0
}
pub fn get_steiner_tree(&self, terminals: &[usize]) -> SteinerTree {
get_steiner_tree(&self.shortest_paths, terminals)
}
}