use std::collections::{HashMap, HashSet, VecDeque}; struct Solution {} impl Solution { fn from_managers(managers: &[i32]) -> HashMap> { let mut mapping = HashMap::new(); for (i, manager) in managers .iter() .enumerate() .filter(|&(_, &manager)| manager != -1) { mapping .entry(*manager as usize) .or_insert_with(HashSet::new) .insert(i); } mapping } pub fn num_of_minutes(_n: i32, head_id: i32, manager: Vec, inform_time: Vec) -> i32 { let manager_mapping = Solution::from_managers(&manager); let mut q = VecDeque::new(); q.push_back((head_id as usize, 0)); let mut minutes = -1; while let Some((manager, t)) = q.pop_front() { minutes = std::cmp::max(minutes, t); for employee in manager_mapping.get(&manager).unwrap_or(&HashSet::new()) { q.push_back((*employee, t + inform_time[manager])); } } minutes } } fn main() {} #[cfg(test)] mod tests { use super::*; // Input: n = 1, headID = 0, manager = [-1], informTime = [0] // Output: 0 // Explanation: The head of the company is the only employee in the company. #[test] fn example_1() { assert_eq!(Solution::num_of_minutes(1, 0, vec![-1], vec![0]), 0); } // Input: n = 6, headID = 2, manager = [2,2,-1,2,2,2], informTime = [0,0,1,0,0,0] // Output: 1 // Explanation: The head of the company with id = 2 is the direct manager of all the employees in the company and needs 1 minute to inform them all. // The tree structure of the employees in the company is shown. #[test] fn example_2() { assert_eq!( Solution::num_of_minutes(6, 2, vec![2, 2, -1, 2, 2, 2], vec![0, 0, 1, 0, 0, 0]), 1 ); } #[test] fn regression_1() { assert_eq!( Solution::num_of_minutes( 15, 0, vec![-1, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6], vec![1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0] ), 3 ); } }