leetcode_rust/
four_sum_2.rs

1#![allow(dead_code)]
2
3// O(n^2): calculus two array's sum firstly and save to hash map, and then eval two sum
4// for other two array. When meet sum1 + sum2 == 0, do counting.
5pub fn four_sum_count(a: Vec<i32>, b: Vec<i32>, c: Vec<i32>, d: Vec<i32>) -> i32 {
6    use std::collections::HashMap;
7
8    let mut sums = HashMap::new();
9    let mut count = 0;
10    for &num_a in a.iter() {
11        for &num_b in b.iter() {
12            let sum = num_a + num_b;
13            match sums.get_mut(&sum) {
14                Some(x) => {
15                    *x += 1;
16                }
17                None => {
18                    sums.insert(sum, 1);
19                }
20            }
21        }
22    }
23
24    for &num_c in c.iter() {
25        for &num_d in d.iter() {
26            let sum = -num_c - num_d;
27            match sums.get(&sum) {
28                Some(x) => {
29                    count += *x;
30                }
31                None => {}
32            }
33        }
34    }
35
36    count
37}
38
39#[cfg(test)]
40mod tests {
41    use super::*;
42
43    #[test]
44    fn test1() {
45        let res = four_sum_count(vec![1, 2], vec![-2, -1], vec![-1, 2], vec![0, 2]);
46        assert_eq!(res, 2);
47        let res = four_sum_count(vec![-1, -1], vec![-1, 1], vec![-1, 1], vec![1, -1]);
48        assert_eq!(res, 6);
49    }
50}
leetcode_rust/four_sum_2.rs

leetcode_rust/
four_sum_2.rs
