leetcode_rust/
three_sum.rs

1#![allow(dead_code)]
2
3// sort and three pointers
4// let target = nums[j]
5// look up i, k when nums[i] + nums[j] = target
6pub fn three_sum(nums: Vec<i32>) -> Vec<Vec<i32>> {
7    if nums.len() < 3 {
8        return vec![];
9    }
10    let mut nums = nums;
11    nums.sort();
12    let mut res = vec![];
13
14    let mut j = nums.len() - 1;
15    while j >= 2 {
16        let mut i = 0;
17        let mut k = j - 1;
18        while i < k {
19            let sum = nums[i] + nums[k];
20            if sum + nums[j] < 0 {
21                i += 1;
22            } else if sum + nums[j] > 0 {
23                k -= 1;
24            } else {
25                res.push(vec![nums[i], nums[k], nums[j]]);
26                loop {
27                    i += 1;
28                    if i >= k || nums[i - 1] != nums[i] {
29                        break;
30                    }
31                }
32                loop {
33                    k -= 1;
34                    if i >= k || nums[k + 1] != nums[k] {
35                        break;
36                    }
37                }
38            }
39        }
40        loop {
41            j -= 1;
42            if j < 2 || nums[j + 1] != nums[j] {
43                break;
44            }
45        }
46    }
47
48    res
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54
55    #[test]
56    fn test1() {
57        let res = three_sum(vec![-1, 0, 1, 2, -1, -4]);
58        assert_eq!(res, vec![vec![-1, -1, 2], vec![-1, 0, 1]]);
59
60        let res = three_sum(vec![0, 1, 2, 3]);
61        let right: Vec<Vec<i32>> = vec![];
62        assert_eq!(res, right);
63
64        let res = three_sum(vec![0, 0, 0]);
65        assert_eq!(res, vec![vec![0, 0, 0]]);
66
67        let res = three_sum(vec![-4, -2, -2, -2, 0, 1, 2, 2, 2, 3, 3, 4, 4, 6, 6]);
68        assert_eq!(
69            res,
70            vec![
71                vec![-4, -2, 6],
72                vec![-4, 0, 4],
73                vec![-2, -2, 4],
74                vec![-4, 1, 3],
75                vec![-4, 2, 2],
76                vec![-2, 0, 2]
77            ]
78        );
79    }
80}