Add problem 3160: Find the Number of Distinct Colors Among the Balls

Author: EFanZh

src/lib.rs

@@ -2113,6 +2113,7 @@ pub mod problem_3152_special_array_ii;
 pub mod problem_3153_sum_of_digit_differences_of_all_pairs;
 pub mod problem_3158_find_the_xor_of_numbers_which_appear_twice;
 pub mod problem_3159_find_occurrences_of_an_element_in_an_array;
+pub mod problem_3160_find_the_number_of_distinct_colors_among_the_balls;
 pub mod problem_3350_adjacent_increasing_subarrays_detection_ii;
 
 #[cfg(test)]

src/problem_3160_find_the_number_of_distinct_colors_among_the_balls/hash_map.rs

@@ -0,0 +1,72 @@
+pub struct Solution;
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+use std::collections::HashMap;
+use std::collections::hash_map::Entry;
+use std::mem;
+
+impl Solution {
+    fn helper(queries: Vec<Vec<i32>>, mut set_color: impl FnMut(i32, i32) -> i32) -> Vec<i32> {
+        let mut color_counts = HashMap::new();
+
+        queries
+            .into_iter()
+            .map(|query| {
+                let [index, color] = query.try_into().ok().unwrap();
+                let old_color = set_color(index, color);
+
+                if color != old_color {
+                    if let Entry::Occupied(entry) = color_counts.entry(old_color) {
+                        if *entry.get() == 0 {
+                            entry.remove();
+                        } else {
+                            *entry.into_mut() -= 1;
+                        }
+                    }
+
+                    match color_counts.entry(color) {
+                        Entry::Occupied(occupied_entry) => *occupied_entry.into_mut() += 1,
+                        Entry::Vacant(vacant_entry) => {
+                            vacant_entry.insert(0);
+                        }
+                    }
+                }
+
+                color_counts.len() as _
+            })
+            .collect()
+    }
+
+    pub fn query_results(limit: i32, queries: Vec<Vec<i32>>) -> Vec<i32> {
+        let limit = limit.cast_unsigned() as usize;
+
+        if limit < queries.len() {
+            let mut colors = vec![0; limit + 1].into_boxed_slice();
+
+            Self::helper(queries, |index, color| {
+                mem::replace(&mut colors[index.cast_unsigned() as usize], color)
+            })
+        } else {
+            let mut colors = HashMap::new();
+
+            Self::helper(queries, |index, color| colors.insert(index, color).unwrap_or(0))
+        }
+    }
+}
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl super::Solution for Solution {
+    fn query_results(limit: i32, queries: Vec<Vec<i32>>) -> Vec<i32> {
+        Self::query_results(limit, queries)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_solution() {
+        super::super::tests::run::<super::Solution>();
+    }
+}

src/problem_3160_find_the_number_of_distinct_colors_among_the_balls/mod.rs

@@ -0,0 +1,24 @@
+pub mod hash_map;
+
+pub trait Solution {
+    fn query_results(limit: i32, queries: Vec<Vec<i32>>) -> Vec<i32>;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Solution;
+
+    pub fn run<S: Solution>() {
+        let test_cases = [
+            ((4, &[[1, 4], [2, 5], [1, 3], [3, 4]] as &[_]), &[1, 2, 2, 3] as &[_]),
+            ((4, &[[0, 1], [1, 2], [2, 2], [3, 4], [4, 5]]), &[1, 2, 2, 3, 4]),
+        ];
+
+        for ((limit, queries), expected) in test_cases {
+            assert_eq!(
+                S::query_results(limit, queries.iter().map(Vec::from).collect()),
+                expected
+            );
+        }
+    }
+}