Improve All Combinations of Size k Implementation

src/backtracking/all_combination_of_size_k.rs

@@ -1,62 +1,123 @@
-/*
-    In this problem, we want to determine all possible combinations of k
-    numbers out of 1 ... n. We use backtracking to solve this problem.
-    Time complexity: O(C(n,k)) which is O(n choose k) = O((n!/(k! * (n - k)!)))
-
-    generate_all_combinations(n=4, k=2) => [[1, 2], [1, 3], [1, 4], [2, 3], [2, 4], [3, 4]]
-*/
-pub fn generate_all_combinations(n: i32, k: i32) -> Vec<Vec<i32>> {
-    let mut result = vec![];
-    create_all_state(1, n, k, &mut vec![], &mut result);
-
-    result
+//! This module provides a function to generate all possible combinations
+//! of `k` numbers out of `0...n-1` using a backtracking algorithm.
+
+/// Custom error type for combination generation.
+#[derive(Debug, PartialEq)]
+pub enum CombinationError {
+    KGreaterThanN,
+    InvalidZeroRange,
+}
+
+/// Generates all possible combinations of `k` numbers out of `0...n-1`.
+///
+/// # Arguments
+///
+/// * `n` - The upper limit of the range (`0` to `n-1`).
+/// * `k` - The number of elements in each combination.
+///
+/// # Returns
+///
+/// A `Result` containing a vector with all possible combinations of `k` numbers out of `0...n-1`,
+/// or a `CombinationError` if the input is invalid.
+pub fn generate_all_combinations(n: usize, k: usize) -> Result<Vec<Vec<usize>>, CombinationError> {
+    if n == 0 && k > 0 {
+        return Err(CombinationError::InvalidZeroRange);
+    }
+
+    if k > n {
+        return Err(CombinationError::KGreaterThanN);
+    }
+
+    let mut combinations = vec![];
+    let mut current = vec![0; k];
+    backtrack(0, n, k, 0, &mut current, &mut combinations);
+    Ok(combinations)
 }
 
-fn create_all_state(
-    increment: i32,
-    total_number: i32,
-    level: i32,
-    current_list: &mut Vec<i32>,
-    total_list: &mut Vec<Vec<i32>>,
+/// Helper function to generate combinations recursively.
+///
+/// # Arguments
+///
+/// * `start` - The current number to start the combination with.
+/// * `n` - The upper limit of the range (`0` to `n-1`).
+/// * `k` - The number of elements left to complete the combination.
+/// * `index` - The current index being filled in the combination.
+/// * `current` - A mutable reference to the current combination being constructed.
+/// * `combinations` - A mutable reference to the vector holding all combinations.
+fn backtrack(
+    start: usize,
+    n: usize,
+    k: usize,
+    index: usize,
+    current: &mut Vec<usize>,
+    combinations: &mut Vec<Vec<usize>>,
 ) {
-    if level == 0 {
-        total_list.push(current_list.clone());
+    if index == k {
+        combinations.push(current.clone());
         return;
     }
 
-    for i in increment..(total_number - level + 2) {
-        current_list.push(i);
-        create_all_state(i + 1, total_number, level - 1, current_list, total_list);
-        current_list.pop();
+    for num in start..=(n - k + index) {
+        current[index] = num;
+        backtrack(num + 1, n, k, index + 1, current, combinations);
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
-    #[test]
-    fn test_output() {
-        let expected_res = vec![
+    macro_rules! combination_tests {
+        ($($name:ident: $test_case:expr,)*) => {
+            $(
+                #[test]
+                fn $name() {
+                    let (n, k, expected) = $test_case;
+                    assert_eq!(generate_all_combinations(n, k), expected);
+                }
+            )*
+        }
+    }
+
+    combination_tests! {
+        test_generate_4_2: (4, 2, Ok(vec![
+            vec![0, 1],
+            vec![0, 2],
+            vec![0, 3],
             vec![1, 2],
             vec![1, 3],
-            vec![1, 4],
             vec![2, 3],
-            vec![2, 4],
-            vec![3, 4],
-        ];
-
-        let res = generate_all_combinations(4, 2);
-
-        assert_eq!(expected_res, res);
-    }
-
-    #[test]
-    fn test_empty() {
-        let expected_res: Vec<Vec<i32>> = vec![vec![]];
-
-        let res = generate_all_combinations(0, 0);
-
-        assert_eq!(expected_res, res);
+        ])),
+        test_generate_4_3: (4, 3, Ok(vec![
+            vec![0, 1, 2],
+            vec![0, 1, 3],
+            vec![0, 2, 3],
+            vec![1, 2, 3],
+        ])),
+        test_generate_5_3: (5, 3, Ok(vec![
+            vec![0, 1, 2],
+            vec![0, 1, 3],
+            vec![0, 1, 4],
+            vec![0, 2, 3],
+            vec![0, 2, 4],
+            vec![0, 3, 4],
+            vec![1, 2, 3],
+            vec![1, 2, 4],
+            vec![1, 3, 4],
+            vec![2, 3, 4],
+        ])),
+        test_generate_5_1: (5, 1, Ok(vec![
+            vec![0],
+            vec![1],
+            vec![2],
+            vec![3],
+            vec![4],
+        ])),
+        test_empty: (0, 0, Ok(vec![vec![]])),
+        test_generate_n_eq_k: (3, 3, Ok(vec![
+            vec![0, 1, 2],
+        ])),
+        test_generate_k_greater_than_n: (3, 4, Err(CombinationError::KGreaterThanN)),
+        test_zero_range_with_nonzero_k: (0, 1, Err(CombinationError::InvalidZeroRange)),
     }
 }