Refactor : 3sum

Author: printjin-gmailcom

3sum/printjin-gmailcom.py

@@ -1,12 +1,14 @@
 from typing import List
+from itertools import combinations
 class Solution:
     def threeSum(self, nums):
+        res = set()
         n = len(nums)
-        result = set()
         for i in range(n):
-            for j in range(i + 1, n):
-                for k in range(j + 1, n):
-                    if nums[i] + nums[j] + nums[k] == 0:
-                        triplet = tuple(sorted([nums[i], nums[j], nums[k]]))
-                        result.add(triplet)
-        return [list(t) for t in result]
+            a = nums[i]
+            rest = nums[:i] + nums[i+1:] 
+            for comb in combinations(rest, 2):
+                if sum(comb) == -a:
+                    triplet = tuple(sorted([a, *comb]))
+                    res.add(triplet)
+        return [list(t) for t in res]