Add problem 2176: Count Equal and Divisible Pairs in an Array

src/lib.rs

@@ -1560,6 +1560,7 @@ pub mod problem_2166_design_bitset;
 pub mod problem_2169_count_operations_to_obtain_zero;
 pub mod problem_2170_minimum_operations_to_make_the_array_alternating;
 pub mod problem_2171_removing_minimum_number_of_magic_beans;
+pub mod problem_2176_count_equal_and_divisible_pairs_in_an_array;
 pub mod problem_2177_find_three_consecutive_integers_that_sum_to_a_given_number;
 pub mod problem_2180_count_integers_with_even_digit_sum;
 pub mod problem_2181_merge_nodes_in_between_zeros;

src/problem_2176_count_equal_and_divisible_pairs_in_an_array/iterative.rs

@@ -0,0 +1,78 @@
+pub struct Solution;
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+use std::cmp::Ordering;
+use std::collections::hash_map::Entry;
+use std::collections::HashMap;
+use std::num::NonZero;
+
+impl Solution {
+    fn gcd(mut x: u8, mut y: NonZero<u8>) -> NonZero<u8> {
+        while let Some(z) = NonZero::<u8>::new(x % y) {
+            x = y.get();
+            y = z;
+        }
+
+        y
+    }
+
+    pub fn count_pairs(nums: Vec<i32>, k: i32) -> i32 {
+        let k = NonZero::<u8>::new(k as _).unwrap();
+        let mut indices = [const { Vec::new() }; 100];
+
+        for (i, num) in (0_u8..).zip(nums) {
+            indices[num as u32 as usize - 1].push(i);
+        }
+
+        // See solutions to <https://leetcode.com/problems/count-array-pairs-divisible-by-k/>.
+
+        let mut counts = HashMap::<_, u16>::new();
+        let mut result = 0;
+
+        for indices in indices {
+            for index in indices {
+                match counts.entry(Self::gcd(index, k)) {
+                    Entry::Occupied(entry) => *entry.into_mut() += 1,
+                    Entry::Vacant(entry) => {
+                        entry.insert(1);
+                    }
+                }
+            }
+
+            for (&gcd_1, &count_1) in &counts {
+                let required = NonZero::<u8>::new(k.get() / gcd_1).unwrap();
+
+                for (&gcd_2, &count_2) in &counts {
+                    if gcd_2.get() % required == 0 {
+                        result += match gcd_1.cmp(&gcd_2) {
+                            Ordering::Less => count_1 * count_2,
+                            Ordering::Equal => count_1 * (count_1 - 1) / 2,
+                            Ordering::Greater => continue,
+                        };
+                    }
+                }
+            }
+
+            counts.clear();
+        }
+
+        u32::from(result) as _
+    }
+}
+
+// ------------------------------------------------------ snip ------------------------------------------------------ //
+
+impl super::Solution for Solution {
+    fn count_pairs(nums: Vec<i32>, k: i32) -> i32 {
+        Self::count_pairs(nums, k)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    #[test]
+    fn test_solution() {
+        super::super::tests::run::<super::Solution>();
+    }
+}

src/problem_2176_count_equal_and_divisible_pairs_in_an_array/mod.rs

@@ -0,0 +1,34 @@
+pub mod iterative;
+
+pub trait Solution {
+    fn count_pairs(nums: Vec<i32>, k: i32) -> i32;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::Solution;
+
+    pub fn run<S: Solution>() {
+        let test_cases = [
+            ((&[3, 1, 2, 2, 2, 1, 3] as &[_], 2), 4),
+            ((&[1, 2, 3, 4], 1), 0),
+            (
+                (
+                    &[
+                        29, 91, 29, 29, 71, 81, 100, 100, 34, 29, 71, 71, 91, 78, 29, 34, 78, 91, 34, 29, 91, 91, 37,
+                        71, 81, 71, 37, 37, 81, 71, 34, 78, 34, 29, 81, 29, 100, 81, 34, 34, 71, 29, 91, 34, 71, 29,
+                        34, 71, 81, 91, 29, 100, 78, 91, 78, 81, 34, 71, 100, 71, 34, 100, 37, 29, 34, 34, 34, 78, 81,
+                        29, 37, 100, 100, 29, 78, 34, 29, 81, 29, 81, 29, 37, 71, 81, 29, 37, 91, 29, 34, 34, 78, 100,
+                        34, 78, 29, 34, 100, 81, 100, 100,
+                    ],
+                    10,
+                ),
+                200,
+            ),
+        ];
+
+        for ((nums, k), expected) in test_cases {
+            assert_eq!(S::count_pairs(nums.to_vec(), k), expected);
+        }
+    }
+}