use std::collections::BinaryHeap;
use std::cmp::Ordering;
#[derive(PartialEq)]
struct NN(f64);
impl Eq for NN {}
impl PartialOrd for NN {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
other.0.partial_cmp(&self.0)
}
}
impl Ord for NN {
fn cmp(&self, other: &Self) -> Ordering {
self.partial_cmp(other).unwrap()
}
}
impl Solution {
pub fn max_probability(n: i32, edges: Vec<Vec<i32>>, succ_prob: Vec<f64>, start_node: i32, end_node: i32) -> f64 {
let n = n as usize;
let mut graph = vec![vec![]; n];
for (edge, sp) in edges.into_iter().zip(succ_prob.into_iter()) {
let (u, v) = (edge[0] as usize, edge[1] as usize);
graph[u].push((v, sp));
graph[v].push((u, sp));
}
let (source, target) = (start_node as usize, end_node as usize);
let mut distances = vec![0.0; n];
distances[source] = 1.0;
let mut queue = BinaryHeap::new();
queue.push((NN(0.0), source));
while let Some((NN(estimate), vertex)) = queue.pop() {
if vertex == target {
return distances[target];
}
if estimate > (1.0 - distances[vertex]) {
continue;
}
for &(neighbour, sp) in &graph[vertex] {
let old_estimate = (1.0 - distances[neighbour]);
let estimate_through_vertex = (1.0 - distances[vertex] * sp);
if estimate_through_vertex < old_estimate {
distances[neighbour] = distances[vertex] * sp;
queue.push((NN(estimate_through_vertex), neighbour));
}
}
}
distances[target]
}
}