Two-sum and Four-sum implemented in Rust

algo.rs

// Actual implementations
pub mod proper {
    use std::collections::{HashSet, HashMap};

    pub fn two_sum(a: &[i32], b: &[i32], n: i32) -> u32 {
        let mut count = 0;
        let mut complements = HashSet::new();

        for x in a.iter() {
            complements.insert(n - x);
        }

        for y in b.iter() {
            if complements.contains(y) {
                count += 1;
            }
        }

        count
    }

    pub fn four_sum(
        a: &[i32],
        b: &[i32],
        c: &[i32],
        d: &[i32],
        n: i32
    ) -> u32 {

        //
        // If      x + y + z + w == n,
        // then    x + y         == n - (z + w) and
        //                 z + w == n - (x + y).
        //

        let mut count = 0;
        let mut complements: HashMap<i32, u32> = HashMap::new();

        // Keep track of all of the 'n - (x + y)'
        for x in a.iter() {
            for y in b.iter() {
                let comp = n - (x + y);
                complements.insert(comp, match complements.get(&comp) {
                    Some(count) => count + 1,
                    None => 1,
                });
            }
        }

        // Check if any of our '(z + w)' match any of the complements
        for z in c.iter() {
            for w in d.iter() {
                let sum = z + w;
                if let Some(count_ab) = complements.get(&sum) {
                    count += count_ab;
                }
            }
        }

        count
    }
}


// Used in unit-testing for comparing against optimized versions, to see if we
// got the right result
pub mod naive {
    pub fn two_sum(a: &[i32], b: &[i32], goal: i32) -> u32 {
        let mut count = 0;

        for x in a.iter() {
            for y in b.iter() {
                if x + y == goal { count += 1; }
            }
        }

        count
    }

    pub fn four_sum(
        a: &[i32],
        b: &[i32],
        c: &[i32],
        d: &[i32],
        goal: i32
    ) -> u32 {
        let mut count = 0;

        for x in a.iter() {
            for y in b.iter() {
                for z in c.iter() {
                    for w in d.iter() {

                        if x + y + z + w == goal { count += 1; }

                    }
                }
            }
        }

        count
    }
}



#[cfg(test)]
mod tests {
    use super::*;
    use rand::Rng;

    static LIST_A: [i32; 26] = [  0,  77, 79, -34,   5, 19, 63, -69, 18,  94, 24, 88,  99,  7, 56, 92,  67, 86,  6, 97, 61, -93, 84, 36, 47, -44 ];
    static LIST_B: [i32; 26] = [ 72,  68, 39,  29,  98, 45, 78,  11, 71,  81, 40, 53, -23, 28, 32, 74,  85, 35, 17, 27,  0,  50, 65, 73, 70,  58 ];
    static LIST_C: [i32; 26] = [ 66,  25, 33,  76, -38, 60, 30,  31,  8,  75, 37, 13,  95, 64, 87, 16,  43,  0, 57, 20, 62,  91, -3, 14, 82,  21 ];
    static LIST_D: [i32; 26] = [ 26, -42, 83,  52,  89, 46,  9,   1, 48, -12, 59, 96,  51, 80,  2, 15, -54,  4, 10, 90, 22,  55,  0, 41, 49,  17 ];

    #[test]
    fn two_sum_test_static() {
        let test_cases = [ 12, 5, 0, 52, 24, 9 ];

        for &goal in test_cases.iter() {
            assert_eq!(
                naive::two_sum(&LIST_A, &LIST_B, goal),
                proper::two_sum(&LIST_A, &LIST_B, goal)
            );

            assert_eq!(
                naive::two_sum(&LIST_C, &LIST_D, goal),
                proper::two_sum(&LIST_C, &LIST_D, goal)
            );
        }
    }

    #[test]
    fn four_sum_test_static() {
        let test_cases = [ 0, 12, 1, 4, 2, 56 ];

        for &goal in test_cases.iter() {
            assert_eq!(
                naive::four_sum(&LIST_A, &LIST_B, &LIST_C, &LIST_D, goal),
                proper::four_sum(&LIST_A, &LIST_B, &LIST_C, &LIST_D, goal)
            );
        }
    }

    fn make_list(n: usize) -> Vec<i32> {
        let mut result = Vec::new();
        for _ in 0..n { result.push(rand::thread_rng().gen_range(-512..512)); }
        result
    }

    #[test]
    fn two_sum_test_random() {
        let a = make_list(50);
        let b = make_list(50);

        let test_cases = make_list(12);

        for &goal in test_cases.iter() {
            assert_eq!(
                naive::two_sum(&a, &b, goal),
                proper::two_sum(&a, &b, goal)
            );
        }
    }

    #[test]
    fn four_sum_test_random() {
        let a = make_list(50);
        let b = make_list(50);
        let c = make_list(50);
        let d = make_list(50);

        let test_cases = make_list(12);

        for &goal in test_cases.iter() {
            assert_eq!(
                naive::four_sum(&a, &b, &c, &d, goal),
                proper::four_sum(&a, &b, &c, &d, goal)
            );
        }
    }
}

$ cargo test
    Finished test [unoptimized + debuginfo] target(s) in 0.03s
     Running unittests (target\debug\deps\algo-98a392387d86ad72.exe)

running 4 tests
test tests::two_sum_test_static ... ok
test tests::two_sum_test_random ... ok
test tests::four_sum_test_static ... ok
test tests::four_sum_test_random ... ok

test result: ok. 4 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 2.99s

   Doc-tests algo

running 0 tests

test result: ok. 0 passed; 0 failed; 0 ignored; 0 measured; 0 filtered out; finished in 0.00s