Google Codejam 2022 Chain Reactions with rust and topological sort

chain-reactions.rs

/*
twitter: @robertkingnz
youtube: https://www.youtube.com/c/RobertKing/videos

topological sort found here: https://github.com/kenkoooo/competitive-programming-rs/blob/master/src/graph/topological_sort.rs
Google Code Jam 2022 Chain Reactions: https://codingcompetitions.withgoogle.com/codejam/round/0000000000876ff1/0000000000a45ef7
*/

pub fn topological_sort(graph: &[Vec<usize>], mut fun: Vec<u64>) -> u64 {
    let n = graph.len();
    let mut min = vec![1_000_000_001u64; n];
    let mut tot = vec![0; n];
    let mut in_deg = vec![0; n];
    for v in 0..n {
        for &to in graph[v].iter() {
            in_deg[to] += 1;
        }
    }

    let mut q = std::collections::VecDeque::new();
    for i in 0..n {
        if in_deg[i] == 0 {
            q.push_back(i);
        }
    }

    let mut ans = 0;
    while let Some(v) = q.pop_front() {
        if graph[v].is_empty() {
            ans += fun[v];
        }
        for &next in graph[v].iter() {
            in_deg[next] -= 1;
            tot[next] += fun[v];
            min[next] = min[next].min(fun[v]);
            if in_deg[next] == 0 {
                ans += tot[next] - min[next];
                fun[next] = fun[next].max(min[next]);
                q.push_back(next);
            }
        }
    }
    ans
}

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 fun: Vec<u64> = sc.vec(n);
        let points: Vec<usize> = sc.vec(n);
        let mut graph = vec![vec![]; n];
        for i in 0..n {
            if points[i] != 0 {
                graph[i].push(points[i]-1);
            }
        }
        let ans = topological_sort(&graph, fun);
        sc.write(
            format!("Case #{}: {}", case_num, ans)
        );
        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()
    }
}