kickstart pizza delivery

pizza_delivery.rs

/*
@robertkingnz
https://www.youtube.com/c/RobertKing/videos -> https://youtu.be/Z5WzJ9jbnKg
https://codingcompetitions.withgoogle.com/kickstart/round/00000000008cb0f5/0000000000ba86e6
*/

extern crate core;

fn calc_toll(c: i64, toll: &(char, i64)) -> i64 {
    let ans = match toll.0 {
        '*' => c * toll.1,
        '/' => ((c as f64) / (toll.1 as f64)).floor() as i64,
        '+' => c + toll.1,
        '-' => c - toll.1,
        _ => panic!("unexpected operator"),
    };
    ans
}

fn neighbours(ar: usize, ac: usize, n: usize, c: i64, tolls: &Vec<(char, i64)>) -> Vec<(usize, usize, i64)> {
    let mut v = vec![(ar, ac, c)];
    if ar > 0 {
        let north = calc_toll(c, &tolls[0]);
        v.push((ar - 1, ac, north));
    }
    if ac > 0 {
        let west = calc_toll(c, &tolls[2]);
        v.push((ar, ac - 1, west));
    }
    if ar + 1 < n {
        let south = calc_toll(c, &tolls[3]);
        v.push((ar + 1, ac, south));
    }
    if ac + 1 < n {
        let east = calc_toll(c, &tolls[1]);
        v.push((ar, ac + 1, east));
    }
    v
}

fn f(n: usize, p: usize, m: usize, ar: usize, ac: usize, tolls: &Vec<(char, i64)>, customer_id: &Vec<Vec<Option<usize>>>, cost: &Vec<i64>) -> Option<i64> {
    // dp t s i j = c
    let mut dp = vec![vec![vec![vec![None; n]; n]; 1 << p]; m+1];
    dp[0][0][ar][ac] = Some(0);
    for t in 0..m {
        for s in 0..1<<p {
            for i in 0..n {
                for j in 0..n {
                    if let Some(c) = dp[t][s][i][j] {
                        for (i2, j2, c2) in neighbours(i, j, n, c, &tolls) {
                            dp[t+1][s][i2][j2] = dp[t+1][s][i2][j2].max(Some(c2));
                            if let Some(idx) = customer_id[i2][j2] {
                                if (s >> idx) & 1 == 0 {
                                    dp[t+1][s | (1 << idx)][i2][j2] = dp[t+1][s | (1 << idx)][i2][j2].max(Some(c2 + cost[idx]));
                                }
                            }
                        }
                    }
                }
            }
        }
    }
    dp[m][(1 << p)-1].iter().flatten().flatten().cloned().max()
}

fn main() {
    assert_eq!((1 as f64/4 as f64).floor() as i64, 0);
    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 {
        // N, P, M, Ar, Ac
        let n: usize = sc.read();
        let p: usize = sc.read();
        let m: usize = sc.read();
        let ar: usize = sc.read();
        let ac: usize = sc.read();
        let mut tolls = vec![];
        for _ in 0..4 {
            let op: char = sc.read();
            let k: i64 = sc.read();
            tolls.push((op, k));
        }
        let mut customer_id = vec![vec![None; n];n];
        let mut cost = vec![0; p];
        for k in 0..p {
            let i: usize = sc.read();
            let j: usize = sc.read();
            let c: i64 = sc.read();
            customer_id[i-1][j-1] = Some(k);
            cost[k] = c;
        }
        let ans = f(n, p, m, ar-1, ac-1, &tolls, &customer_id, &cost);
        if let Some(ans) = ans {
            sc.write(
                format!("Case #{}: {}", case_num, ans)
            );
        } else {
            sc.write(
                format!("Case #{}: {}", case_num, "IMPOSSIBLE")
            );
        }
        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
    }
}