evaluate_division.rs

use std::collections::{btree_map::Values, HashMap, HashSet};

struct Solution;
impl Solution {
    pub fn calc_equation(
        equations: Vec<Vec<String>>,
        values: Vec<f64>,
        queries: Vec<Vec<String>>,
    ) -> Vec<f64> {
        let mut adj: HashMap<&str, Vec<(&str, f64)>> = HashMap::new();
        for (v, val) in equations.iter().zip(values.iter()) {
            adj.entry(v[0].as_str())
                .or_default()
                .push((v[1].as_str(), *val));
            adj.entry(v[1].as_str())
                .or_default()
                .push((v[0].as_str(), 1.0 / val));
        }

        fn calculate<'a>(
            adj: &HashMap<&str, Vec<(&'a str, f64)>>,
            vs: &mut HashSet<&'a str>,
            target: &str,
            s: &'a str,
        ) -> (bool, f64) {
            if vs.contains(s) {
                return (false, 1.0);
            }
            if s == target {
                return (true, 1.0);
            }
            vs.insert(s);

            if adj.contains_key(s) {
                for (x, y) in adj.get(s).unwrap().iter() {
                    let (done, value) = calculate(adj, vs, target, x);
                    if done {
                        return (true, value * y);
                    }
                }
            }

            (false, 1.0)
        }
        let mut anss = vec![];
        for v in queries.iter() {
            let (a, b) = (v[0].as_str(), v[1].as_str());
            let mut vs = HashSet::new();
            let (done, ans) = calculate(&adj, &mut vs, b, a);
            println!("{} {}: {:?} {}", a, b, done, ans);
            if done && adj.contains_key(a) && adj.contains_key(b) {
                anss.push(ans);
            } else {
                anss.push(-1.0);
            }
        }
        anss
    }
}

fn main() {
    let equations = vec![vec!["a".into(), "b".into()], vec!["b".into(), "c".into()]];
    let values = [2.0, 3.0];
    let queries = vec![
        vec!["a".into(), "c".into()],
        vec!["b".into(), "a".into()],
        vec!["a".into(), "e".into()],
        vec!["a".into(), "a".into()],
        vec!["x".into(), "x".into()],
    ];
    println!(
        "{:?}",
        Solution::calc_equation(equations, values.into(), queries)
    );
}