From 43c3f4ea4070bfbe1f41f4b861c7ff3f89953715 Mon Sep 17 00:00:00 2001 From: Bama Charan Chhandogi Date: Sun, 1 Oct 2023 20:16:12 +0530 Subject: [PATCH] add Three sum (#9177) * add Three sum * add Three sum * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update * update * update * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add documention --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- maths/three_sum.py | 47 ++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) create mode 100644 maths/three_sum.py diff --git a/maths/three_sum.py b/maths/three_sum.py new file mode 100644 index 000000000..09956f841 --- /dev/null +++ b/maths/three_sum.py @@ -0,0 +1,47 @@ +""" +https://en.wikipedia.org/wiki/3SUM +""" + + +def three_sum(nums: list[int]) -> list[list[int]]: + """ + Find all unique triplets in a sorted array of integers that sum up to zero. + + Args: + nums: A sorted list of integers. + + Returns: + A list of lists containing unique triplets that sum up to zero. + + >>> three_sum([-1, 0, 1, 2, -1, -4]) + [[-1, -1, 2], [-1, 0, 1]] + >>> three_sum([1, 2, 3, 4]) + [] + """ + nums.sort() + ans = [] + for i in range(len(nums) - 2): + if i == 0 or (nums[i] != nums[i - 1]): + low, high, c = i + 1, len(nums) - 1, 0 - nums[i] + while low < high: + if nums[low] + nums[high] == c: + ans.append([nums[i], nums[low], nums[high]]) + + while low < high and nums[low] == nums[low + 1]: + low += 1 + while low < high and nums[high] == nums[high - 1]: + high -= 1 + + low += 1 + high -= 1 + elif nums[low] + nums[high] < c: + low += 1 + else: + high -= 1 + return ans + + +if __name__ == "__main__": + import doctest + + doctest.testmod()