LeetCode/rs/detonate-the-maximum-bombs.rs

use std::collections::{HashMap, HashSet, VecDeque};

fn distance(a: (i64, i64), b: (i64, i64)) -> i64 {
    let (x1, y1) = a;
    let (x2, y2) = b;

    (x1 - x2).pow(2) + (y1 - y2).pow(2)
}

struct Graph {
    edges: HashMap<usize, HashSet<usize>>,
}

impl Graph {
    fn new(bombs: &[Vec<i32>]) -> Self {
        let mut edges = HashMap::new();

        for (i, bomb) in bombs.iter().enumerate() {
            edges.insert(i, HashSet::new());
            let neighbours = edges.get_mut(&i).unwrap();

            let (x, y, r) = (bomb[0], bomb[1], bomb[2]);
            let rr = (r as i64).pow(2);

            for j in 0..bombs.len() {
                if i == j {
                    continue;
                }

                let (x1, y1) = (bombs[j][0], bombs[j][1]);

                if distance((x as i64, y as i64), (x1 as i64, y1 as i64)) <= rr {
                    neighbours.insert(j);
                }
            }
        }

        Self { edges }
    }

    fn bfs(&self, mut visited: Vec<bool>, u0: usize) -> i32 {
        let mut found = 0;

        let mut q = VecDeque::new();
        q.push_back(u0);
        visited[u0] = true;

        while let Some(u) = q.pop_front() {
            found += 1;

            for &v in self.edges.get(&u).unwrap_or(&HashSet::new()) {
                if !visited[v] {
                    q.push_back(v);
                    visited[v] = true;
                }
            }
        }

        found
    }
}

struct Solution {}
impl Solution {
    pub fn maximum_detonation(bombs: Vec<Vec<i32>>) -> i32 {
        let graph = Graph::new(&bombs);
        let visited = vec![false; bombs.len()];
        (0..bombs.len())
            .map(|u| graph.bfs(visited.clone(), u))
            .max()
            .unwrap_or(0)
    }
}

fn main() {}

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

    #[test]
    fn example_1() {
        assert_eq!(
            Solution::maximum_detonation(vec![vec![2, 1, 3], vec![6, 1, 4]]),
            2
        );
    }

    #[test]
    fn example_2() {
        assert_eq!(
            Solution::maximum_detonation(vec![vec![1, 1, 5], vec![10, 10, 5]]),
            1
        );
    }

    #[test]
    fn example_3() {
        assert_eq!(
            Solution::maximum_detonation(vec![
                vec![1, 2, 3],
                vec![2, 3, 1],
                vec![3, 4, 2],
                vec![4, 5, 3],
                vec![5, 6, 4]
            ]),
            5
        );
    }

    #[test]
    fn regression_1() {
        assert_eq!(
            Solution::maximum_detonation(vec![vec![1, 1, 100000], vec![100000, 100000, 1]]),
            1
        );
    }
}