diff --git a/src/problem/mod.rs b/src/problem/mod.rs index 863cef1..489a876 100644 --- a/src/problem/mod.rs +++ b/src/problem/mod.rs @@ -609,3 +609,5 @@ mod p2176_count_equal_and_divisible_pairs_in_an_array; mod p2364_count_number_of_bad_pairs; mod p2563_count_the_number_of_fair_pairs; + +mod p781_rabbits_in_forest; diff --git a/src/problem/p2563_count_the_number_of_fair_pairs.rs b/src/problem/p2563_count_the_number_of_fair_pairs.rs index edcb0d0..22cc2c1 100644 --- a/src/problem/p2563_count_the_number_of_fair_pairs.rs +++ b/src/problem/p2563_count_the_number_of_fair_pairs.rs @@ -14,7 +14,7 @@ impl Solution { let l = nums[0..i].partition_point(|&x| x + nums[i] < lower); let r = nums[0..i].partition_point(|&x| x + nums[i] <= upper); - result += (r - l); + result += r - l; } result as i64 diff --git a/src/problem/p781_rabbits_in_forest.rs b/src/problem/p781_rabbits_in_forest.rs new file mode 100644 index 0000000..e521e3e --- /dev/null +++ b/src/problem/p781_rabbits_in_forest.rs @@ -0,0 +1,46 @@ +/** + * [781] Rabbits in Forest + */ +pub struct Solution {} + +// submission codes start here +use std::collections::HashMap; + +impl Solution { + pub fn num_rabbits(answers: Vec) -> i32 { + let mut map = HashMap::new(); + + for i in answers.into_iter() { + let entry = map.entry(i).or_insert(0); + *entry += 1; + } + + let mut result = 0; + + for (k, mut v) in map.into_iter() { + // 同种颜色兔子最多有k + 1个 + let max_count = k + 1; + + while v > 0 { + result += max_count; + v -= max_count; + } + } + + result + } +} + +// submission codes end + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_781() { + assert_eq!(5, Solution::num_rabbits(vec![1, 0, 1, 0, 0])); + assert_eq!(5, Solution::num_rabbits(vec![1, 1, 2])); + assert_eq!(11, Solution::num_rabbits(vec![10, 10, 10])) + } +}