LeetCode/problems/pacific-atlantic-water-flow.rs

use std::{collections::VecDeque, ops::BitOr};

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(u8)]
enum Reachable {
    None = 0x0,
    Pacific = 0x1,
    Atlantic = 0x2,
    Both = 0x3,
}

impl Reachable {
    fn new(y: usize, x: usize, max_y: usize, max_x: usize) -> Reachable {
        let reaches_pacific = y == 0 || x == 0;
        let reaches_atlantic = y == max_y - 1 || x == max_x - 1;

        if reaches_atlantic && reaches_pacific {
            Reachable::Both
        } else if reaches_atlantic {
            Reachable::Atlantic
        } else if reaches_pacific {
            Reachable::Pacific
        } else {
            Reachable::None
        }
    }
}

impl From<u8> for Reachable {
    fn from(u: u8) -> Self {
        match u {
            0 => Reachable::None,
            1 => Reachable::Pacific,
            2 => Reachable::Atlantic,
            3 => Reachable::Both,
            _ => panic!("Invalid conversion"),
        }
    }
}

impl From<Reachable> for u8 {
    fn from(r: Reachable) -> Self {
        match r {
            Reachable::None => 0,
            Reachable::Pacific => 1,
            Reachable::Atlantic => 2,
            Reachable::Both => 3,
        }
    }
}

impl BitOr for Reachable {
    type Output = Reachable;

    fn bitor(self, rhs: Self) -> Self::Output {
        let (lhs, rhs): (u8, u8) = (self.into(), rhs.into());
        (lhs | rhs).into()
    }
}

#[derive(PartialEq, Eq)]
enum State {
    Unvisited,
    Queued,
    Visited,
}

struct Vertex {
    height: i32,
    reaches: Reachable,
    state: State,
}

struct Graph {
    vertices: Vec<Vec<Vertex>>,
    max_y: usize,
    max_x: usize,
}

impl Graph {
    pub fn new(heights: &Vec<Vec<i32>>) -> Graph {
        let max_y = heights.len();
        let max_x = heights[0].len();

        Graph {
            vertices: heights
                .iter()
                .enumerate()
                .map(|(y, row)| {
                    row.iter()
                        .enumerate()
                        .map(|(x, h)| Vertex {
                            height: *h,
                            reaches: Reachable::new(y, x, max_y, max_x),
                            state: State::Unvisited,
                        })
                        .collect()
                })
                .collect(),
            max_y,
            max_x,
        }
    }

    pub fn vertices(&self) -> &Vec<Vec<Vertex>> {
        &self.vertices
    }

    fn _vertex(&mut self, y: i32, x: i32) -> Option<&mut Vertex> {
        if y >= 0 && y < self.max_y as i32 && x >= 0 && x < self.max_x as i32 {
            Some(&mut self.vertices[y as usize][x as usize])
        } else {
            None
        }
    }

    fn _run_bfs(&mut self, q: &mut VecDeque<(i32, i32)>) {
        q.iter().for_each(|&(y, x)| {
            self._vertex(y, x).unwrap().state = State::Queued;
        });

        while let Some((y, x)) = q.pop_front() {
            self._vertex(y, x).unwrap().state = State::Visited;
            let h = self._vertex(y, x).unwrap().height;
            let r = self._vertex(y, x).unwrap().reaches;

            for d in [-1, 1] {
                for (next_y, next_x) in [(y + d, x), (y, x + d)] {
                    if let Some(v) = self._vertex(next_y, next_x) {
                        if matches!(v.state, State::Unvisited) && h <= v.height {
                            q.push_back((next_y, next_x));
                            v.state = State::Queued;
                            v.reaches = r | v.reaches;
                        }
                    }
                }
            }
        }
    }

    fn _reset(&mut self) {
        for y in 0..self.max_y {
            for x in 0..self.max_x {
                self.vertices[y][x].state = State::Unvisited;
            }
        }
    }

    pub fn bfs(&mut self) {
        // Run BFS from Pacific Ocean
        self._run_bfs(&mut VecDeque::from(
            (0..self.max_x)
                .map(|x| (0_i32, x as i32))
                .chain((0..self.max_y).map(|y| (y as i32, 0_i32)))
                .collect::<Vec<(i32, i32)>>(),
        ));

        // Reset the visited
        self._reset();

        // Run BFS from Atlantic Ocean
        self._run_bfs(&mut VecDeque::from(
            (0..self.max_x)
                .map(|x| (self.max_y as i32 - 1, x as i32))
                .chain((0..self.max_y).map(|y| (y as i32, self.max_x as i32 - 1)))
                .collect::<Vec<(i32, i32)>>(),
        ));
    }
}

impl Solution {
    pub fn pacific_atlantic(heights: Vec<Vec<i32>>) -> Vec<Vec<i32>> {
        let mut g = Graph::new(&heights);
        g.bfs();

        g.vertices().iter().for_each(|row| {
            row.iter().for_each(|v| {
                print!("{:?} ", v.reaches);
            });

            println!();
        });

        g.vertices()
            .iter()
            .enumerate()
            .flat_map(|(y, row)| {
                row.iter()
                    .enumerate()
                    .filter(|&(_, v)| matches!(v.reaches, Reachable::Both))
                    .map(move |(x, _)| vec![y as i32, x as i32])
            })
            .collect()
    }
}
struct Solution {}

fn main() {
    println!("Hello World!");
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_example_1() {
        assert_eq!(
            Solution::pacific_atlantic(vec![
                vec![1, 2, 2, 3, 5],
                vec![3, 2, 3, 4, 4],
                vec![2, 4, 5, 3, 1],
                vec![6, 7, 1, 4, 5],
                vec![5, 1, 1, 2, 4]
            ]),
            vec![
                vec![0, 4],
                vec![1, 3],
                vec![1, 4],
                vec![2, 2],
                vec![3, 0],
                vec![3, 1],
                vec![4, 0]
            ]
        );
    }

    #[test]
    fn test_example_2() {
        assert_eq!(Solution::pacific_atlantic(vec![vec![1]]), vec![vec![0, 0]]);
    }

    #[test]
    fn test_wrong() {
        assert_eq!(
            Solution::pacific_atlantic(vec![vec![3, 3, 3], vec![3, 1, 3], vec![0, 2, 4]]),
            vec![
                vec![0, 0],
                vec![0, 1],
                vec![0, 2],
                vec![1, 0],
                vec![1, 2],
                vec![2, 0],
                vec![2, 1],
                vec![2, 2]
            ]
        );
    }
}