diff --git a/src/bin/15_3sum.rs b/src/bin/15_3sum.rs index f5c0f2b..5494883 100644 --- a/src/bin/15_3sum.rs +++ b/src/bin/15_3sum.rs @@ -1,26 +1,32 @@ fn main() { println!("{:?}", three_sum(vec![-1,0,1,2,-1,-4])); + println!("{:?}", three_sum(vec![0,0,0,0,0,0])); + println!("{:?}", three_sum(vec![-2,0,0,2,2])); } // TODO: fix and optimize pub fn three_sum(mut nums: Vec) -> Vec> { - let mut result: Vec> = Vec::new(); - let mut seen = std::collections::HashSet::::with_capacity(nums.len()); - let mut seen2 = std::collections::HashSet::::with_capacity(nums.len()-1); + let mut result = Vec::>::new(); nums.sort(); for i in 0..nums.len()-2 { - if seen.contains(&nums[i]) { continue; } - for j in i+1..nums.len()-1 { - if seen2.contains(&nums[j]) { continue; } - for k in j+1..nums.len() { - if nums[i] + nums[j] + nums[k] == 0 { + if i > 0 && nums[i] == nums[i - 1] { continue; } + let (mut j, mut k) = (i + 1, nums.len() - 1); + while j < k { + if j > i + 1 && nums[j - 1] == nums[j] { + j += 1; continue; + } + if k < nums.len() - 1 && nums[k + 1] == nums[k] { + k -= 1; continue; + } + match (nums[i] + nums[j] + nums[k]).cmp(&0) { + std::cmp::Ordering::Less => { j += 1; }, + std::cmp::Ordering::Greater => { k -= 1 }, + std::cmp::Ordering::Equal => { result.push(vec![nums[i], nums[j], nums[k]]); + j += 1; k -= 1; } } - seen2.insert(nums[j]); } - seen.insert(nums[i]); - seen2.clear(); } return result; }