leetcode/src/problem/p399_evaluate_division.rs
2024-06-23 11:19:08 +08:00

123 lines
2.8 KiB
Rust

/**
* [399] Evaluate Division
*/
pub struct Solution {}
// submission codes start here
use std::collections::HashMap;
struct UnionFind {
parents : Vec<usize>,
weights : Vec<f64>
}
impl UnionFind {
fn new(n: usize) -> UnionFind {
let mut parents = Vec::with_capacity(n);
let mut weights = Vec::with_capacity(n);
for i in 0..n {
parents.push(i);
weights.push(1f64);
}
UnionFind {
parents,
weights
}
}
fn find(&mut self, x: usize) -> usize {
if self.parents[x] == x {
return x;
}
let origin = self.parents[x];
self.parents[x] = self.find(origin);
self.weights[x] *= self.weights[origin];
self.parents[x]
}
fn merge(&mut self, x: usize, y: usize, value: f64) {
let (x_root, y_root) = (self.find(x), self.find(y));
if x_root == y_root {
return;
}
self.parents[x_root] = y_root;
self.weights[x_root] = value * self.weights[y] / self.weights[x];
}
fn is_connected(&mut self, x: usize, y: usize) -> f64 {
let (x_root, y_root) = (self.find(x), self.find(y));
if x_root == y_root {
return self.weights[x] / self.weights[y];
}
-1f64
}
}
impl Solution {
pub fn calc_equation(equations: Vec<Vec<String>>, values: Vec<f64>, queries: Vec<Vec<String>>) -> Vec<f64> {
let count = equations.len();
let mut union_find = UnionFind::new(count * 2);
let mut id_map = HashMap::with_capacity(count * 2);
let mut id = 0;
for (i, equation) in equations.iter().enumerate() {
let (x, y) = (&equation[0], &equation[1]);
let x_id = *id_map.entry(x).or_insert_with(|| {
let inserted_id = id;
id += 1;
inserted_id
});
let y_id = *id_map.entry(y).or_insert_with(|| {
let inserted_id = id;
id += 1;
inserted_id
});
union_find.merge(x_id, y_id, values[i]);
}
let mut result = Vec::with_capacity(queries.len());
for query in &queries {
let (x, y) = (&query[0], &query[1]);
let (x_id, y_id) = (id_map.get(x), id_map.get(y));
if x_id.is_none() || y_id.is_none() {
result.push(-1f64);
} else {
result.push(union_find.is_connected(*x_id.unwrap(),*y_id.unwrap()));
}
}
result
}
}
// submission codes end
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_399() {
let mut map = HashMap::new();
let x = "a".to_owned();
let y = "a".to_owned();
map.entry(&x).or_insert(1);
assert_eq!(map.get(&y), Some(&1));
}
}