Add problem 2261: K Divisible Elements Subarrays

src/lib.rs

@@ -1611,6 +1611,7 @@ pub mod problem_2255_count_prefixes_of_a_given_string;
 pub mod problem_2256_minimum_average_difference;
 pub mod problem_2259_remove_digit_from_number_to_maximize_result;
 pub mod problem_2260_minimum_consecutive_cards_to_pick_up;
+pub mod problem_2261_k_divisible_elements_subarrays;
 pub mod problem_2262_total_appeal_of_a_string;
 pub mod problem_2264_largest_3_same_digit_number_in_string;
 pub mod problem_2265_count_nodes_equal_to_average_of_subtree;

src/problem_2261_k_divisible_elements_subarrays/mod.rs

@@ -0,0 +1,18 @@
+pub mod sliding_window_and_hash_set;
+
+pub trait Solution {
+    fn count_distinct(nums: Vec<i32>, k: i32, p: i32) -> i32;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Solution;
+
+    pub fn run<S: Solution>() {
+        let test_cases = [((&[2, 3, 3, 2, 2] as &[_], 2, 2), 11), ((&[1, 2, 3, 4], 4, 1), 10)];
+
+        for ((nums, k, p), expected) in test_cases {
+            assert_eq!(S::count_distinct(nums.to_vec(), k, p), expected);
+        }
+    }
+}

src/problem_2261_k_divisible_elements_subarrays/sliding_window_and_hash_set.rs

@@ -0,0 +1,48 @@
+pub struct Solution;
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+use std::collections::HashSet;
+use std::num::NonZeroU8;
+
+impl Solution {
+    pub fn count_distinct(nums: Vec<i32>, k: i32, p: i32) -> i32 {
+        let k = k as u8;
+        let p = NonZeroU8::new(p as _).unwrap();
+        let nums = nums.into_iter().map(|num| num as u8).collect::<Vec<_>>();
+        let mut start = 0;
+        let mut count = 0;
+        let mut sub_arrays = HashSet::new();
+
+        for (i, &num) in (1..).zip(&nums) {
+            count += u8::from(num % p == 0);
+
+            while count > k {
+                count -= u8::from(nums[start] % p == 0);
+                start += 1;
+            }
+
+            for start in start..i {
+                sub_arrays.insert(&nums[start..i]);
+            }
+        }
+
+        sub_arrays.len() as _
+    }
+}
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl super::Solution for Solution {
+    fn count_distinct(nums: Vec<i32>, k: i32, p: i32) -> i32 {
+        Self::count_distinct(nums, k, p)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_solution() {
+        super::super::tests::run::<super::Solution>();
+    }
+}