meta hacker cup: Problem D: Second Flight

second-flight.rs

// @robertkingnz
// https://www.youtube.com/c/RobertKing/videos
// https://www.facebook.com/codingcompetitions/hacker-cup/2022/qualification-round/problems/D

use std::collections::{HashMap};
use std::mem;

fn main() {
    let (r, w) = (std::io::stdin(), std::io::stdout());
    let mut sc = IO::new(r.lock(), w.lock());

    let t: usize = sc.read();
    for case_num in 1..=t {
        let n: usize = sc.read();
        let m: usize = sc.read();
        let q: usize = sc.read();
        eprintln!("case {}, {} {} {}", case_num, n, m, q);
        let mut g = vec![HashMap::new(); n];
        let heavy: usize = (n.max(m).max(q) as f64).sqrt() as usize;
        for _ in 0..m {
            let mut ai: usize = sc.read();
            ai -= 1;
            let mut bi: usize = sc.read();
            bi -= 1;
            let ci: u64 = sc.read();
            g[ai].insert(bi, ci);
            g[bi].insert(ai, ci);
        }
        let mut queries = Vec::with_capacity(q);
        for _ in 0..q {
            let mut xi: usize = sc.read();
            xi -= 1;
            let mut yi: usize = sc.read();
            yi -= 1;
            queries.push((xi, yi));
        }
        let mut answers = vec![0; q];

        let mut cache = HashMap::new();
        // cache all queries involving heavy nodes, each one in O(N). and there's only Sqrt(N) heavy nodes.
        for mut x in 0..n {
            if g[x].len() >= heavy {
                for mut y in 0..n {
                    if y != x {
                        if g[x].len() < g[y].len() {
                            mem::swap(&mut x, &mut y);
                        }
                        let mut tot = *g[x].get(&y).unwrap_or(&0) * 2;
                        for (&z, &c1) in &g[y] {
                            let c2 = *g[z].get(&x).unwrap_or(&0);
                            tot += c1.min(c2);
                        }
                        cache.insert((x, y), tot);
                        cache.insert((y, x), tot);
                    }
                }
            }
        }
        // solve remaining light light queries
        for i in 0..q {
            let (x, y) = queries[i];
            if let Some(&tot) = cache.get(&(x, y)) {
                answers[i] = tot;
                continue;
            }
            answers[i] = *g[x].get(&y).unwrap_or(&0) * 2;
            for (&z, &c1) in &g[x] {
                let c2 = *g[z].get(&y).unwrap_or(&0);
                answers[i] += c1.min(c2);
            }
        }
        let ans = answers.iter().map(|tot| format!("{}", tot)).collect::<Vec<String>>();
        sc.write(
            format!("Case #{}: {}", case_num, ans.join(" "))
        );
        sc.write("\n");
    }
}

pub struct IO<R, W: std::io::Write>(R, std::io::BufWriter<W>);

impl<R: std::io::Read, W: std::io::Write> IO<R, W> {
    pub fn new(r: R, w: W) -> IO<R, W> {
        IO(r, std::io::BufWriter::new(w))
    }
    pub fn write<S: ToString>(&mut self, s: S) {
        use std::io::Write;
        self.1.write_all(s.to_string().as_bytes()).unwrap();
    }
    pub fn read<T: std::str::FromStr>(&mut self) -> T {
        use std::io::Read;
        let buf = self
            .0
            .by_ref()
            .bytes()
            .map(|b| b.unwrap())
            .skip_while(|&b| b == b' ' || b == b'\n' || b == b'\r' || b == b'\t')
            .take_while(|&b| b != b' ' && b != b'\n' && b != b'\r' && b != b'\t')
            .collect::<Vec<_>>();
        unsafe { std::str::from_utf8_unchecked(&buf) }
            .parse()
            .ok()
            .expect("Parse error.")
    }
    pub fn usize0(&mut self) -> usize {
        self.read::<usize>() - 1
    }
    pub fn vec<T: std::str::FromStr>(&mut self, n: usize) -> Vec<T> {
        (0..n).map(|_| self.read()).collect()
    }
    pub fn chars(&mut self) -> Vec<char> {
        self.read::<String>().chars().collect()
    }

    pub fn binary_vec(&mut self) -> Vec<u8> {
        self.read::<String>()
            .bytes()
            .map(|b| b - b'0')
            .collect()
    }

    pub fn grid<T, F>(&mut self, r: usize, f: &F) -> Vec<Vec<T>>
        where F: Fn(char) -> T {
        let mut g = Vec::with_capacity(r);
        for _ in 0..r {
            g.push(
                self.chars().into_iter().map(f).collect()
            )
        }
        g
    }
}