count_lakes.rs

use std::collections::VecDeque;

type Point = (i32, i32);

fn count_number_of_lakes(grid: Vec<Vec<usize>>) -> usize {
    let mut count = 0;
    let n = grid.len();
    let m = grid[0].len();
    if n == 0 || m == 0 {
        return 0;
    }
    let mut visited = vec![vec![false; m]; n];

    let directions: Vec<Point> = vec![
        (0, 1),  // North
        (0, -1), // South
        (-1, 0), // West
        (1, 0),  // East
    ];

    for i in 0..n {
        for j in 0..m {
            if grid[i][j] == 0 && !visited[i][j] {
                // use explicit stack to avoid blowing up the call stack
                let mut stack = VecDeque::new();
                stack.push_back((i, j));
                visited[i][j] = true;
                while let Some((x, y)) = stack.pop_back() {
                    for (a, b) in directions.iter() {
                        let new_x = (x as i32 + *a) as usize;
                        let new_y = (y as i32 + *b) as usize;
                        if new_x < n
                            && new_y < m
                            && grid[new_x][new_y] == 0
                            && !visited[new_x][new_y]
                        {
                            visited[new_x][new_y] = true;
                            stack.push_back((new_x, new_y));
                        }
                    }
                }
                count += 1;
            }
        }
    }
    count
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn case_1() {
        let grid = vec![vec![0, 0, 0],
                        vec![1, 0, 1],
                        vec![1, 0, 0]];
        let result = count_number_of_lakes(grid);
        assert_eq!(1, result);
    }

    #[test]
    fn case_2() {
        let grid = vec![vec![0, 0, 0, 1],
                        vec![1, 0, 1, 0],
                        vec![1, 0, 0, 1]];
        let result = count_number_of_lakes(grid);
        assert_eq!(2, result);
    }

    #[test]
    fn case_3() {
        let grid = vec![vec![0, 0],
                        vec![1, 0]];

        let result = count_number_of_lakes(grid);
        assert_eq!(1, result);
    }

    #[test]
    fn case_4() {
        let grid = vec![vec![1, 1],
                        vec![1, 1]];

        let result = count_number_of_lakes(grid);
        assert_eq!(0, result);
    }

    #[test]
    fn case_5() {
        let grid = vec![vec![1, 1, 0, 0],
                        vec![0, 1, 1, 1],
                        vec![0, 1, 0, 1],
                        vec![0, 1, 0, 1],
                        vec![1, 1, 1, 1],
                        vec![0, 1, 0, 1]];

        let result = count_number_of_lakes(grid);
        assert_eq!(5, result);
    }

    #[test]
    fn case_6() {
        let grid = vec![vec![]];

        let result = count_number_of_lakes(grid);
        assert_eq!(0, result);
    }
}