78 lines
2.2 KiB
Rust
78 lines
2.2 KiB
Rust
use std::collections::{HashMap, HashSet, VecDeque};
|
|
|
|
struct Solution {}
|
|
impl Solution {
|
|
fn from_managers(managers: &[i32]) -> HashMap<usize, HashSet<usize>> {
|
|
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<i32>, inform_time: Vec<i32>) -> 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
|
|
);
|
|
}
|
|
}
|